Add alpha matting node (#2292)

chaiNNer-org · Oct 31, 2023 · dff3c12 · dff3c12
1 parent 6da28a6
commit dff3c12
Show file tree

Hide file tree

Showing 13 changed files with 120 additions and 1,106 deletions.
diff --git a/backend/src/nodes/impl/rembg/bg.py b/backend/src/nodes/impl/rembg/bg.py
@@ -16,13 +16,11 @@
 )
 from PIL import Image
 from PIL.Image import Image as PILImage
+from pymatting import estimate_alpha_cf, estimate_foreground_ml
 from scipy.ndimage import binary_erosion
 
 from ...impl.image_utils import normalize
 from ...utils.utils import get_h_w_c
-from .pymatting.estimate_alpha_cf import estimate_alpha_cf
-from .pymatting.estimate_foreground_ml import estimate_foreground_ml
-from .pymatting.util import stack_images
 from .session_factory import new_session
 
 kernel = getStructuringElement(MORPH_ELLIPSE, (3, 3))
@@ -62,7 +60,7 @@ def alpha_matting_cutout(
 
     alpha = estimate_alpha_cf(img_normalized, trimap_normalized)
     foreground = estimate_foreground_ml(img_normalized, alpha)
-    cutout = stack_images(foreground, alpha)
+    cutout = np.dstack((foreground, alpha))
 
     cutout = np.clip(cutout * 255, 0, 255).astype(np.uint8)  # type: ignore
     cutout = Image.fromarray(cutout)

diff --git a/backend/src/nodes/impl/rembg/pymatting/LICENSE.md b/backend/src/nodes/impl/rembg/pymatting/LICENSE.md
diff --git a/backend/src/nodes/impl/rembg/pymatting/__init__.py b/backend/src/nodes/impl/rembg/pymatting/__init__.py
diff --git a/backend/src/nodes/impl/rembg/pymatting/cf_laplacian.py b/backend/src/nodes/impl/rembg/pymatting/cf_laplacian.py
diff --git a/backend/src/nodes/impl/rembg/pymatting/cg.py b/backend/src/nodes/impl/rembg/pymatting/cg.py