Upload Salia_Croppytools.py

Salia_Croppytools.py

@@ -3,6 +3,9 @@ from typing import Tuple
 
 import torch
 import torch.nn.functional as F
+import numpy as np
+from PIL import Image
+
 
 # Salia utils (same style as your loader node)
 try:
@@ -28,7 +31,7 @@ except Exception:
 
 
 # -----------------------------
-# Helpers
+# Helpers (IMAGE)
 # -----------------------------
 
 def _as_image(img: torch.Tensor) -> torch.Tensor:
@@ -42,48 +45,6 @@ def _as_image(img: torch.Tensor) -> torch.Tensor:
     return img
 
 
-def _as_mask(msk: torch.Tensor) -> torch.Tensor:
-    # ComfyUI MASK is usually [B,H,W] float 0..1
-    if not isinstance(msk, torch.Tensor):
-        raise TypeError("MASK must be a torch.Tensor")
-    if msk.dim() == 2:
-        msk = msk.unsqueeze(0)
-    if msk.dim() != 3:
-        raise ValueError(f"Expected MASK shape [B,H,W] (or [H,W]), got {tuple(msk.shape)}")
-    return msk
-
-
-def _match_batch(a: torch.Tensor, b: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-    ba = a.shape[0]
-    bb = b.shape[0]
-    if ba == bb:
-        return a, b
-    if ba == 1 and bb > 1:
-        return a.expand(bb, *a.shape[1:]), b
-    if bb == 1 and ba > 1:
-        return a, b.expand(ba, *b.shape[1:])
-    raise ValueError(f"Batch mismatch: A has batch {ba}, B has batch {bb} (and neither is 1).")
-
-
-def _resize_mask_to(msk: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor:
-    # msk: [B,H,W] -> resize to [B,target_h,target_w]
-    if msk.shape[1] == target_h and msk.shape[2] == target_w:
-        return msk
-    x = msk.unsqueeze(1)  # [B,1,H,W]
-    x = F.interpolate(x, size=(target_h, target_w), mode="bilinear", align_corners=False)
-    return x.squeeze(1)
-
-
-def _combine_alpha_union(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-    """
-    "Alpha combine" (union) like standard alpha coverage:
-        out = 1 - (1-a)*(1-b)
-    """
-    a = a.clamp(0.0, 1.0)
-    b = b.clamp(0.0, 1.0)
-    return (1.0 - (1.0 - a) * (1.0 - b)).clamp(0.0, 1.0)
-
-
 def _crop_with_padding(image: torch.Tensor, x: int, y: int, w: int, h: int) -> torch.Tensor:
     """
     Crops [x,y] top-left, size w*h. If out of bounds, pads with zeros.
@@ -141,10 +102,17 @@ def _alpha_over_region(overlay: torch.Tensor, canvas: torch.Tensor, x: int, y: i
     overlay = _as_image(overlay)
     canvas = _as_image(canvas)
 
-    overlay, canvas = _match_batch(overlay, canvas)
+    # Simple batch handling (Comfy usually matches batches, but allow 1->N)
+    if overlay.shape[0] != canvas.shape[0]:
+        if overlay.shape[0] == 1 and canvas.shape[0] > 1:
+            overlay = overlay.expand(canvas.shape[0], *overlay.shape[1:])
+        elif canvas.shape[0] == 1 and overlay.shape[0] > 1:
+            canvas = canvas.expand(overlay.shape[0], *canvas.shape[1:])
+        else:
+            raise ValueError(f"Batch mismatch: overlay {overlay.shape[0]} vs canvas {canvas.shape[0]}")
 
     B, Hc, Wc, Cc = canvas.shape
-    Bo, Ho, Wo, Co = overlay.shape
+    _, Ho, Wo, _ = overlay.shape
 
     x = int(x)
     y = int(y)
@@ -199,6 +167,65 @@ def _alpha_over_region(overlay: torch.Tensor, canvas: torch.Tensor, x: int, y: i
     return out
 
 
+# -----------------------------
+# RMBG EXACT MASK COMBINE LOGIC (copied solution)
+# -----------------------------
+
+class _AILab_MaskCombiner_Exact:
+    def combine_masks(self, mask_1, mode="combine", mask_2=None, mask_3=None, mask_4=None):
+        try:
+            masks = [m for m in [mask_1, mask_2, mask_3, mask_4] if m is not None]
+
+            if len(masks) <= 1:
+                return (masks[0] if masks else torch.zeros((1, 64, 64), dtype=torch.float32),)
+
+            ref_shape = masks[0].shape
+            masks = [self._resize_if_needed(m, ref_shape) for m in masks]
+
+            if mode == "combine":
+                result = torch.maximum(masks[0], masks[1])
+                for mask in masks[2:]:
+                    result = torch.maximum(result, mask)
+            elif mode == "intersection":
+                result = torch.minimum(masks[0], masks[1])
+            else:
+                result = torch.abs(masks[0] - masks[1])
+
+            return (torch.clamp(result, 0, 1),)
+        except Exception as e:
+            print(f"Error in combine_masks: {str(e)}")
+            print(f"Mask shapes: {[m.shape for m in masks]}")
+            raise e
+
+    def _resize_if_needed(self, mask, target_shape):
+        try:
+            if mask.shape == target_shape:
+                return mask
+
+            if len(mask.shape) == 2:
+                mask = mask.unsqueeze(0)
+            elif len(mask.shape) == 4:
+                mask = mask.squeeze(1)
+
+            target_height = target_shape[-2] if len(target_shape) >= 2 else target_shape[0]
+            target_width = target_shape[-1] if len(target_shape) >= 2 else target_shape[1]
+
+            resized_masks = []
+            for i in range(mask.shape[0]):
+                mask_np = mask[i].cpu().numpy()
+                img = Image.fromarray((mask_np * 255).astype(np.uint8))
+                img_resized = img.resize((target_width, target_height), Image.LANCZOS)
+                mask_resized = np.array(img_resized).astype(np.float32) / 255.0
+                resized_masks.append(torch.from_numpy(mask_resized))
+
+            return torch.stack(resized_masks)
+
+        except Exception as e:
+            print(f"Error in _resize_if_needed: {str(e)}")
+            print(f"Input mask shape: {mask.shape}, Target shape: {target_shape}")
+            raise e
+
+
 # -----------------------------
 # 1) Cropout_Square_From_IMG
 # -----------------------------
@@ -282,7 +309,7 @@ class Paste_rect_to_img:
 
 
 # -----------------------------
-# 4) Combine_2_masks
+# 4) Combine_2_masks  (RMBG exact: torch.maximum + PIL resize)
 # -----------------------------
 
 class Combine_2_masks:
@@ -290,30 +317,20 @@ class Combine_2_masks:
 
     @classmethod
     def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "maskA": ("MASK",),
-                "maskB": ("MASK",),
-            }
-        }
+        return {"required": {"maskA": ("MASK",), "maskB": ("MASK",)}}
 
     RETURN_TYPES = ("MASK",)
     RETURN_NAMES = ("mask",)
     FUNCTION = "run"
 
     def run(self, maskA, maskB):
-        a = _as_mask(maskA)
-        b = _as_mask(maskB)
-
-        a, b = _match_batch(a, b)
-        b = _resize_mask_to(b, a.shape[1], a.shape[2])
-
-        out = _combine_alpha_union(a, b)
+        combiner = _AILab_MaskCombiner_Exact()
+        out, = combiner.combine_masks(maskA, mode="combine", mask_2=maskB)
         return (out,)
 
 
 # -----------------------------
-# 5) Combine_2_masks_invert_1
+# 5) Combine_2_masks_invert_1 (invert A then RMBG combine)
 # -----------------------------
 
 class Combine_2_masks_invert_1:
@@ -321,31 +338,22 @@ class Combine_2_masks_invert_1:
 
     @classmethod
     def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "maskA": ("MASK",),
-                "maskB": ("MASK",),
-            }
-        }
+        return {"required": {"maskA": ("MASK",), "maskB": ("MASK",)}}
 
     RETURN_TYPES = ("MASK",)
     RETURN_NAMES = ("mask",)
     FUNCTION = "run"
 
     def run(self, maskA, maskB):
-        a = _as_mask(maskA)
-        b = _as_mask(maskB)
-
-        a, b = _match_batch(a, b)
-        b = _resize_mask_to(b, a.shape[1], a.shape[2])
-
-        a_inv = (1.0 - a).clamp(0.0, 1.0)
-        out = _combine_alpha_union(a_inv, b)
+        combiner = _AILab_MaskCombiner_Exact()
+        maskA = 1.0 - maskA
+        out, = combiner.combine_masks(maskA, mode="combine", mask_2=maskB)
         return (out,)
 
 
 # -----------------------------
 # 6) Combine_2_masks_inverse
+# invert both, combine, invert result (RMBG max logic)
 # -----------------------------
 
 class Combine_2_masks_inverse:
@@ -353,34 +361,24 @@ class Combine_2_masks_inverse:
 
     @classmethod
     def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "maskA": ("MASK",),
-                "maskB": ("MASK",),
-            }
-        }
+        return {"required": {"maskA": ("MASK",), "maskB": ("MASK",)}}
 
     RETURN_TYPES = ("MASK",)
     RETURN_NAMES = ("mask",)
     FUNCTION = "run"
 
     def run(self, maskA, maskB):
-        a = _as_mask(maskA)
-        b = _as_mask(maskB)
-
-        a, b = _match_batch(a, b)
-        b = _resize_mask_to(b, a.shape[1], a.shape[2])
-
-        a_inv = (1.0 - a).clamp(0.0, 1.0)
-        b_inv = (1.0 - b).clamp(0.0, 1.0)
-
-        combined_inv = _combine_alpha_union(a_inv, b_inv)
-        out = (1.0 - combined_inv).clamp(0.0, 1.0)  # == a*b (intersection)
+        combiner = _AILab_MaskCombiner_Exact()
+        maskA = 1.0 - maskA
+        maskB = 1.0 - maskB
+        combined, = combiner.combine_masks(maskA, mode="combine", mask_2=maskB)
+        out = 1.0 - combined
+        out = torch.clamp(out, 0, 1)
         return (out,)
 
 
 # -----------------------------
-# 7) combine_masks_with_loaded
+# 7) combine_masks_with_loaded (RMBG exact combine)
 # -----------------------------
 
 class combine_masks_with_loaded:
@@ -404,16 +402,10 @@ class combine_masks_with_loaded:
         if image == "<no pngs found>":
             raise FileNotFoundError("No PNGs in assets/images")
 
-        base = _as_mask(mask)
-
-        # Load image+mask from assets (Salia util)
         _img, loaded_mask = load_image_from_assets(image)
-        loaded = _as_mask(loaded_mask)
 
-        base, loaded = _match_batch(base, loaded)
-        loaded = _resize_mask_to(loaded, base.shape[1], base.shape[2])
-
-        out = _combine_alpha_union(base, loaded)
+        combiner = _AILab_MaskCombiner_Exact()
+        out, = combiner.combine_masks(mask, mode="combine", mask_2=loaded_mask)
         return (out,)
 
     @classmethod
@@ -435,6 +427,105 @@ class combine_masks_with_loaded:
         return True
 
 
+# -----------------------------
+# 8) NEW: invert input mask, combine with loaded mask, apply to image alpha, paste on canvas
+# -----------------------------
+
+class apply_segment:
+    CATEGORY = "image/salia"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        choices = list_pngs() or ["<no pngs found>"]
+        return {
+            "required": {
+                "mask": ("MASK",),
+                "image": (choices, {}),          # dropdown asset (used ONLY for loaded mask)
+                "img": ("IMAGE",),               # the image to receive final_mask as alpha (overlay source)
+                "canvas": ("IMAGE",),            # destination
+                "x": ("INT", {"default": 0, "min": -100000, "max": 100000, "step": 1}),
+                "y": ("INT", {"default": 0, "min": -100000, "max": 100000, "step": 1}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = "run"
+
+    def run(self, mask, image, img, canvas, x, y):
+        if image == "<no pngs found>":
+            raise FileNotFoundError("No PNGs in assets/images")
+
+        combiner = _AILab_MaskCombiner_Exact()
+
+        # Load asset mask (do NOT invert)
+        _img_asset, loaded_mask = load_image_from_assets(image)
+
+        # Invert input mask, then combine with loaded mask (RMBG exact combine => maximum)
+        inv_mask = 1.0 - mask
+        final_mask, = combiner.combine_masks(inv_mask, mode="combine", mask_2=loaded_mask)
+
+        # Apply final_mask as alpha to input image -> final_overlay (RGBA)
+        img = _as_image(img)
+        B, H, W, C = img.shape
+
+        # Resize final_mask to match img H/W if needed (uses RMBG exact resize helper)
+        # (target_shape must look like a mask shape [B,H,W], but resize keeps its own batch count)
+        final_mask_resized = combiner._resize_if_needed(final_mask, (final_mask.shape[0], H, W))
+
+        # Batch match (simple 1->N expansion only)
+        if final_mask_resized.shape[0] != B:
+            if final_mask_resized.shape[0] == 1 and B > 1:
+                final_mask_resized = final_mask_resized.expand(B, H, W)
+            elif B == 1 and final_mask_resized.shape[0] > 1:
+                img = img.expand(final_mask_resized.shape[0], *img.shape[1:])
+                B = img.shape[0]
+            else:
+                raise ValueError(f"Batch mismatch: img batch={B}, final_mask batch={final_mask_resized.shape[0]}")
+
+        if C == 3:
+            # RGB -> RGBA with alpha = final_mask
+            alpha = final_mask_resized.to(device=img.device, dtype=img.dtype)
+            final_overlay = torch.cat([img, alpha.unsqueeze(-1)], dim=-1)
+        else:
+            # RGBA: combine existing alpha with final_mask using RMBG combine (maximum)
+            rgb = img[..., :3]
+            alpha_img = img[..., 3]  # [B,H,W]
+
+            # RMBG combine uses PIL-resize sometimes, so keep combine inputs on CPU
+            a1 = alpha_img.detach().cpu()
+            a2 = final_mask_resized.detach().cpu()
+            combined_alpha, = combiner.combine_masks(a1, mode="combine", mask_2=a2)
+
+            combined_alpha = combined_alpha.to(device=img.device, dtype=img.dtype)
+            final_overlay = torch.cat([rgb, combined_alpha.unsqueeze(-1)], dim=-1)
+
+        # Paste final_overlay onto canvas at (x,y)
+        canvas = _as_image(canvas)
+        final_overlay = final_overlay.to(device=canvas.device, dtype=canvas.dtype)
+
+        out = _alpha_over_region(final_overlay, canvas, x, y)
+        return (out,)
+
+    @classmethod
+    def IS_CHANGED(cls, mask, image, img, canvas, x, y):
+        if image == "<no pngs found>":
+            return image
+        return file_hash(image)
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, mask, image, img, canvas, x, y):
+        if image == "<no pngs found>":
+            return "No PNGs in assets/images"
+        try:
+            path = safe_path(image)
+        except Exception as e:
+            return str(e)
+        if not os.path.isfile(path):
+            return f"File not found in assets/images: {image}"
+        return True
+
+
 # -----------------------------
 # Node mappings
 # -----------------------------
@@ -447,6 +538,7 @@ NODE_CLASS_MAPPINGS = {
     "Combine_2_masks_invert_1": Combine_2_masks_invert_1,
     "Combine_2_masks_inverse": Combine_2_masks_inverse,
     "combine_masks_with_loaded": combine_masks_with_loaded,
+    "apply_segment": apply_segment,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -457,4 +549,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "Combine_2_masks_invert_1": "Combine_2_masks_invert_1",
     "Combine_2_masks_inverse": "Combine_2_masks_inverse",
     "combine_masks_with_loaded": "combine_masks_with_loaded",
+    "apply_segment": "apply_segment",
 }