Changes made

Author: Code2Collapse

nodes/mask_preview.py

@@ -97,8 +97,8 @@ def preview(self, image, mask, display_mode, overlay_color_r, overlay_color_g,
                 overlay = color.view(1, 1, 3).expand(H, W, 3)
                 frame = frame * (1 - edge_3d) + overlay * edge_3d
 
-            # Edge overlay for non-edge modes
-            if display_mode != "edge_highlight" and edge_width > 0:
+            # Edge overlay for non-edge modes (skip side_by_side — frame width differs)
+            if display_mode not in ("edge_highlight", "side_by_side") and edge_width > 0:
                 edge = self._detect_edge(mi, edge_width)
                 edge_3d = edge.unsqueeze(-1)
                 overlay_edge = color.view(1, 1, 3).expand(H, W, 3)

nodes/sam_mask_generator.py

@@ -9,11 +9,14 @@
 import json
 import gc
 
-try:
-    import cv2
-    HAS_CV2 = True
-except ImportError:
-    HAS_CV2 = False
+from .utils import (
+    HAS_CV2,
+    get_sam_predictor,
+    augment_prompts_from_mask,
+    mask_to_sam_logits,
+    parse_points_json,
+    parse_bbox_input,
+)
 
 
 class SAMMaskGeneratorMEC:
@@ -127,42 +130,20 @@ def _run_inference(self, model, model_type, image, points_json, bbox_json,
         img_np = (img_tensor.cpu().numpy() * 255).astype(np.uint8)
         H, W = img_np.shape[:2]
 
-        # Parse points
-        try:
-            points_list = json.loads(points_json) if isinstance(points_json, str) else points_json
-        except json.JSONDecodeError:
-            points_list = []
-
-        point_coords = None
-        point_labels = None
+        # Parse points (shared utility)
+        points_list = parse_points_json(points_json)
+        point_coords, point_labels = None, None
         if points_list:
             coords = [[float(p["x"]), float(p["y"])] for p in points_list]
             labels = [int(p.get("label", 1)) for p in points_list]
             point_coords = np.array(coords, dtype=np.float32)
             point_labels = np.array(labels, dtype=np.int32)
 
-        # Parse bbox
-        box_np = None
-        if bbox_input is not None:
-            # From BBOX node: [x, y, w, h] → [x1, y1, x2, y2]
-            bx, by, bw, bh = bbox_input
-            box_np = np.array([bx, by, bx + bw, by + bh], dtype=np.float32)
-        elif bbox_json and bbox_json.strip():
-            try:
-                bdata = json.loads(bbox_json)
-                if isinstance(bdata, list) and len(bdata) == 4:
-                    box_np = np.array(bdata, dtype=np.float32)
-                elif isinstance(bdata, dict):
-                    bx = float(bdata.get("x", 0))
-                    by = float(bdata.get("y", 0))
-                    bw = float(bdata.get("w", bdata.get("width", 0)))
-                    bh = float(bdata.get("h", bdata.get("height", 0)))
-                    box_np = np.array([bx, by, bx + bw, by + bh], dtype=np.float32)
-            except json.JSONDecodeError:
-                pass
-
-        # ── Build predictor ────────────────────────────────────────────
-        predictor = self._get_predictor(model, model_type, img_np)
+        # Parse bbox (shared utility)
+        box_np = parse_bbox_input(bbox_json, bbox_input)
+
+        # ── Build predictor (shared utility) ───────────────────────────
+        predictor = get_sam_predictor(model, model_type, img_np)
 
         if predictor is None:
             empty = torch.zeros(1, H, W, dtype=torch.float32)
@@ -188,17 +169,17 @@ def _run_inference(self, model, model_type, image, points_json, bbox_json,
             iter_labels = point_labels
             iter_box = box_np
 
-            # Augment prompts from previous iteration
+            # Augment prompts from previous iteration (shared utility)
             if iteration > 0 and current_mask is not None:
-                iter_coords, iter_labels, iter_box = self._augment_prompts(
+                iter_coords, iter_labels, iter_box = augment_prompts_from_mask(
                     current_mask, point_coords, point_labels, box_np,
-                    auto_negative_points, H, W,
+                    H, W, auto_negative=auto_negative_points,
                 )
 
             # Prepare mask input (logits from previous pass)
             mask_input = None
             if iteration > 0 and current_mask is not None:
-                mask_input = self._mask_to_logits(current_mask)
+                mask_input = mask_to_sam_logits(current_mask)
 
             # Run SAM
             try:
@@ -296,107 +277,6 @@ def _run_inference(self, model, model_type, image, points_json, bbox_json,
 
         return (selected_mask, all_masks_t, det_bbox, selected_score, info)
 
-    # ── Predictor factory ─────────────────────────────────────────────
-
-    def _get_predictor(self, model, model_type, img_np):
-        """Get the correct predictor for the model type and set image."""
-        predictor = None
-
-        if model_type in ("sam2", "sam2.1"):
-            try:
-                from sam2.sam2_image_predictor import SAM2ImagePredictor
-                predictor = SAM2ImagePredictor(model)
-            except Exception:
-                pass
-
-        elif model_type == "sam3":
-            try:
-                from sam3.predictor import SAM3Predictor
-                predictor = SAM3Predictor(model)
-            except Exception:
-                pass
-
-        else:
-            try:
-                from segment_anything import SamPredictor
-                predictor = SamPredictor(model)
-            except Exception:
-                pass
-
-        if predictor is not None:
-            try:
-                predictor.set_image(img_np)
-            except Exception:
-                return None
-
-        return predictor
-
-    # ── Iterative prompt augmentation ─────────────────────────────────
-
-    def _augment_prompts(self, mask, orig_coords, orig_labels, orig_box,
-                          auto_neg, H, W):
-        """Generate augmented prompts from the previous mask pass."""
-        coords_list = []
-        labels_list = []
-
-        if orig_coords is not None:
-            coords_list.append(orig_coords)
-            labels_list.append(orig_labels)
-
-        if not HAS_CV2:
-            c = np.concatenate(coords_list) if coords_list else None
-            l = np.concatenate(labels_list) if labels_list else None
-            return c, l, orig_box
-
-        binary = (mask > 0.5).astype(np.uint8)
-
-        # Interior positive points
-        kern = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
-        interior = cv2.erode(binary, kern, iterations=1)
-        pts = np.argwhere(interior > 0)
-        if len(pts) > 0:
-            n = min(3, len(pts))
-            idx = np.linspace(0, len(pts) - 1, n, dtype=int)
-            coords_list.append(pts[idx][:, ::-1].astype(np.float32))
-            labels_list.append(np.ones(n, dtype=np.int32))
-
-        # Exterior negative points
-        if auto_neg:
-            dilated = cv2.dilate(binary, kern, iterations=1)
-            exterior = dilated - binary
-            pts = np.argwhere(exterior > 0)
-            if len(pts) > 0:
-                n = min(2, len(pts))
-                idx = np.linspace(0, len(pts) - 1, n, dtype=int)
-                coords_list.append(pts[idx][:, ::-1].astype(np.float32))
-                labels_list.append(np.zeros(n, dtype=np.int32))
-
-        all_coords = np.concatenate(coords_list).astype(np.float32) if coords_list else None
-        all_labels = np.concatenate(labels_list).astype(np.int32) if labels_list else None
-
-        # Derive tighter bbox
-        ys, xs = np.where(binary > 0)
-        if len(xs) > 0:
-            pad = max(5, int(min(H, W) * 0.02))
-            box = np.array([
-                max(0, xs.min() - pad), max(0, ys.min() - pad),
-                min(W, xs.max() + pad), min(H, ys.max() + pad),
-            ], dtype=np.float32)
-        else:
-            box = orig_box
-
-        return all_coords, all_labels, box
-
-    @staticmethod
-    def _mask_to_logits(mask, target_size=256):
-        """Convert float mask → SAM logit space (inverse sigmoid) at 256x256."""
-        m = np.clip(mask, 1e-6, 1 - 1e-6)
-        logits = np.log(m / (1 - m))
-        if HAS_CV2:
-            logits = cv2.resize(logits, (target_size, target_size),
-                                 interpolation=cv2.INTER_LINEAR)
-        return logits[np.newaxis, :, :]  # (1, 256, 256)
-
     @staticmethod
     def _fallback_predict(model, img_np, point_coords, point_labels, box_np, multimask):
         """Generic fallback using model's forward pass directly."""

nodes/sam_vitmatte_pipeline.py

@@ -36,6 +36,11 @@
     remove_small_regions,
     mask_to_bbox,
     make_split_preview,
+    augment_prompts_from_mask,
+    mask_to_sam_logits,
+    parse_points_json,
+    points_to_arrays,
+    parse_bbox_input,
 )
 
 try:
@@ -147,9 +152,9 @@ def execute(self, sam_model, image, points_json, bbox_json,
         img_np = (img_tensor.cpu().numpy() * 255).astype(np.uint8)
         H, W = img_np.shape[:2]
 
-        # Parse prompts
-        points_list = self._parse_points(points_json)
-        box_np = self._parse_bbox(bbox_json, bbox)
+        # Parse prompts (shared utilities)
+        points_list = parse_points_json(points_json)
+        box_np = parse_bbox_input(bbox_json, bbox)
 
         # ── Stage 1: SAM coarse mask (with iterative refinement) ──────
         if offload and hasattr(model, "to"):
@@ -194,7 +199,7 @@ def execute(self, sam_model, image, points_json, bbox_json,
         # ── Outputs ───────────────────────────────────────────────────
         edge_mask = torch.abs(refined_mask - (coarse_mask > 0.5).float())
         det_bbox = mask_to_bbox(refined_mask, W, H)
-        preview = make_split_preview(img_tensor, refined_mask, coarse_mask)
+        preview = make_split_preview(img_tensor, coarse_mask, refined_mask)
 
         info = json.dumps({
             "model_type": model_type,
@@ -232,7 +237,7 @@ def _iterative_sam(self, model, model_type, img_np, points_list,
         if predictor is None:
             return torch.zeros((H, W), dtype=torch.float32), 0.0
 
-        point_coords, point_labels = self._points_to_arrays(points_list)
+        point_coords, point_labels = points_to_arrays(points_list)
 
         # Use existing mask as starting point if provided
         current_mask = None
@@ -251,7 +256,7 @@ def _iterative_sam(self, model, model_type, img_np, points_list,
             iter_box = box_np
 
             if iteration > 0 and current_mask is not None:
-                aug_coords, aug_labels, aug_box = self._augment_prompts_from_mask(
+                aug_coords, aug_labels, aug_box = augment_prompts_from_mask(
                     current_mask, point_coords, point_labels, box_np, H, W
                 )
                 iter_coords = aug_coords
@@ -263,14 +268,7 @@ def _iterative_sam(self, model, model_type, img_np, points_list,
             try:
                 mask_input = None
                 if iteration > 0 and current_mask is not None:
-                    m_logit = np.clip(current_mask, 1e-6, 1 - 1e-6)
-                    m_logit = np.log(m_logit / (1 - m_logit))
-                    if HAS_CV2:
-                        mask_input = cv2.resize(m_logit, (256, 256),
-                                                 interpolation=cv2.INTER_LINEAR)
-                    else:
-                        mask_input = m_logit
-                    mask_input = mask_input[np.newaxis, :, :]  # (1, 256, 256)
+                    mask_input = mask_to_sam_logits(current_mask)
 
                 masks_np, scores, _ = predictor.predict(
                     point_coords=iter_coords,
@@ -310,63 +308,6 @@ def _iterative_sam(self, model, model_type, img_np, points_list,
 
         return torch.from_numpy(current_mask), best_score
 
-    def _augment_prompts_from_mask(self, mask, orig_coords, orig_labels,
-                                    orig_box, H, W):
-        """Generate additional prompts from the previous mask iteration."""
-        coords_list = []
-        labels_list = []
-
-        if orig_coords is not None:
-            coords_list.append(orig_coords)
-            labels_list.append(orig_labels)
-
-        if not HAS_CV2:
-            c = np.concatenate(coords_list) if coords_list else None
-            l = np.concatenate(labels_list) if labels_list else None
-            return c, l, orig_box
-
-        binary = (mask > 0.5).astype(np.uint8)
-
-        # Eroded interior → strong positive points
-        kern = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
-        interior = cv2.erode(binary, kern, iterations=1)
-        interior_pts = np.argwhere(interior > 0)  # (y, x)
-        if len(interior_pts) > 0:
-            n = min(3, len(interior_pts))
-            indices = np.linspace(0, len(interior_pts) - 1, n, dtype=int)
-            sampled = interior_pts[indices]
-            coords_list.append(sampled[:, ::-1].astype(np.float32))
-            labels_list.append(np.ones(n, dtype=np.int32))
-
-        # Dilated boundary exterior → negative points
-        dilated = cv2.dilate(binary, kern, iterations=1)
-        exterior = dilated - binary
-        exterior_pts = np.argwhere(exterior > 0)
-        if len(exterior_pts) > 0:
-            n = min(2, len(exterior_pts))
-            indices = np.linspace(0, len(exterior_pts) - 1, n, dtype=int)
-            sampled = exterior_pts[indices]
-            coords_list.append(sampled[:, ::-1].astype(np.float32))
-            labels_list.append(np.zeros(n, dtype=np.int32))
-
-        all_coords = np.concatenate(coords_list).astype(np.float32) if coords_list else None
-        all_labels = np.concatenate(labels_list).astype(np.int32) if labels_list else None
-
-        # Derive tighter bbox from mask
-        ys, xs = np.where(binary > 0)
-        if len(xs) > 0:
-            pad = max(5, int(min(H, W) * 0.02))
-            box = np.array([
-                max(0, xs.min() - pad),
-                max(0, ys.min() - pad),
-                min(W, xs.max() + pad),
-                min(H, ys.max() + pad),
-            ], dtype=np.float32)
-        else:
-            box = orig_box
-
-        return all_coords, all_labels, box
-
     # ══════════════════════════════════════════════════════════════════
     #  STAGE 3 – Edge-aware matting (delegates to utils)
     # ══════════════════════════════════════════════════════════════════
@@ -456,42 +397,3 @@ def _try_laplacian(mask_np, edge_radius, detail_pres):
             return torch.from_numpy(np.clip(result, 0, 1).astype(np.float32))
         except Exception:
             return None
-
-    # ══════════════════════════════════════════════════════════════════
-    #  Prompt parsing
-    # ══════════════════════════════════════════════════════════════════
-
-    @staticmethod
-    def _parse_points(points_json):
-        try:
-            return json.loads(points_json) if isinstance(points_json, str) else points_json
-        except json.JSONDecodeError:
-            return []
-
-    @staticmethod
-    def _points_to_arrays(points_list):
-        if not points_list:
-            return None, None
-        coords = [[float(p["x"]), float(p["y"])] for p in points_list]
-        labels = [int(p.get("label", 1)) for p in points_list]
-        return np.array(coords, dtype=np.float32), np.array(labels, dtype=np.int32)
-
-    @staticmethod
-    def _parse_bbox(bbox_json, bbox_input):
-        if bbox_input is not None:
-            bx, by, bw, bh = bbox_input
-            return np.array([bx, by, bx + bw, by + bh], dtype=np.float32)
-        if bbox_json and bbox_json.strip():
-            try:
-                bdata = json.loads(bbox_json)
-                if isinstance(bdata, list) and len(bdata) == 4:
-                    return np.array(bdata, dtype=np.float32)
-                elif isinstance(bdata, dict):
-                    bx = float(bdata.get("x", 0))
-                    by = float(bdata.get("y", 0))
-                    bw = float(bdata.get("w", bdata.get("width", 0)))
-                    bh = float(bdata.get("h", bdata.get("height", 0)))
-                    return np.array([bx, by, bx + bw, by + bh], dtype=np.float32)
-            except json.JSONDecodeError:
-                pass
-        return None