Sam3 instance bonding

meshroom/sam3/ImageSegmentationSam3.py

@@ -51,6 +51,19 @@ class ImageSegmentationSam3(desc.Node):
             description="Weights file for the segmentation model.",
             value="${RDS_SAM3_MODEL_PATH}",
         ),
+        desc.BoolParam(
+            name="enableBonding",
+            label="Enable Masks Bonding",
+            description="Enable bonding where instances overlap.",
+            value=True,
+        ),
+        desc.IntParam(
+            name="bondingKernelSize",
+            label="Bonding Kernel Size",
+            description="Kernel size for morphological processing applied for masks bonding.",
+            value=11,
+            enabled=lambda node: node.enableBonding.value,
+        ),
         desc.BoolParam(
             name="maskInvert",
             label="Invert Masks",
@@ -238,11 +251,17 @@ def updateDetectedBboxes(self, detectedBBoxes, xyxy, idx, key, color):
                 "height": xyxy[3] - xyxy[1]
                 }}
 
-    def updateMaskImageAndDetectedBboxes(self, inference_state, maskImage, detectedBBoxes, key, w_ori, h_ori, PAR, orientation):
-        from segmentationRDS import image
+    def updateMaskImageAndDetectedBboxes(self, inference_state, maskImage, detectedBBoxes, key, w_ori, h_ori, PAR, orientation, ks_bond = 0):
+        from segmentationRDS import image, sam3Utils
+        import numpy as np
         masks, boxes, scores = inference_state["masks"], inference_state["boxes"], inference_state["scores"]
-        for mask in masks:
-            maskImage[mask.squeeze(0).cpu()] = [255, 255, 255]
+        masks = [mask.squeeze(0).cpu().numpy() for mask in masks]
+        if masks:
+            masks_stack = np.stack(masks, axis=0)
+            bool_mask = np.sum(masks_stack, axis=0) > 0
+            if len(masks) > 1 and ks_bond > 0:
+                bool_mask = sam3Utils.bond_masks(masks, ks_bond, ks_bond, ks_bond) > 0
+            maskImage[bool_mask] = [255, 255, 255]
         for idx, box in enumerate(boxes):
             x1, y1, x2, y2 = box.cpu().tolist()
             x1, y1 = image.fromUsualToRawOrientation(x1, y1, w_ori, h_ori, PAR, orientation)
@@ -349,21 +368,24 @@ def processChunk(self, chunk):
                         inference_state = processor.set_text_prompt(state=inference_state, prompt=textPrompt)
                         # Get the masks, bounding boxes, and scores
                         if "masks" in inference_state:
-                            self.updateMaskImageAndDetectedBboxes(inference_state, mask_image, detectedShapeBboxes, key, w_ori, h_ori, PAR, orientation)
+                            ks_bond = chunk.node.bondingKernelSize.value if chunk.node.enableBonding.value else 0
+                            self.updateMaskImageAndDetectedBboxes(inference_state, mask_image, detectedShapeBboxes, key, w_ori, h_ori, PAR, orientation, ks_bond)
 
-                    if not chunk.node.splitBoxPrompt:
+                    if not chunk.node.splitBoxPrompt.value:
                         processor.reset_all_prompts(state=inference_state)
                     for box, label in zip(bboxes, bboxLabels):
                         # Prompt the model with bboxes
-                        if chunk.node.splitBoxPrompt:
+                        if chunk.node.splitBoxPrompt.value:
                             processor.reset_all_prompts(state=inference_state)
                         inference_state = processor.add_geometric_prompt(state=inference_state, box=box, label=label)
                         # Get the masks, bounding boxes, and scores
-                        if "masks" in inference_state and label and chunk.node.splitBoxPrompt:
-                            self.updateMaskImageAndDetectedBboxes(inference_state, mask_image, detectedShapeBboxes, key, w_ori, h_ori, PAR, orientation)
+                        if "masks" in inference_state and label and chunk.node.splitBoxPrompt.value:
+                            ks_bond = chunk.node.bondingKernelSize.value if chunk.node.enableBonding.value else 0
+                            self.updateMaskImageAndDetectedBboxes(inference_state, mask_image, detectedShapeBboxes, key, w_ori, h_ori, PAR, orientation, ks_bond)
 
-                    if "masks" in inference_state and not chunk.node.splitBoxPrompt:
-                        self.updateMaskImageAndDetectedBboxes(inference_state, mask_image, detectedShapeBboxes, key, w_ori, h_ori, PAR, orientation)
+                    if "masks" in inference_state and not chunk.node.splitBoxPrompt.value:
+                        ks_bond = chunk.node.bondingKernelSize.value if chunk.node.enableBonding.value else 0
+                        self.updateMaskImageAndDetectedBboxes(inference_state, mask_image, detectedShapeBboxes, key, w_ori, h_ori, PAR, orientation, ks_bond)
 
                     if chunk.node.maskInvert.value:
                         mask = (mask_image[:,:,0:1] == 0).astype('float32')

meshroom/sam3/VideoSegmentationSam3.py

@@ -62,6 +62,19 @@ class VideoSegmentationSam3(desc.Node):
             value=16,
             enabled=lambda node: node.timeSlicing.value,
         ),
+        desc.BoolParam(
+            name="enableBonding",
+            label="Enable Masks Bonding",
+            description="Enable bonding where instances overlap.",
+            value=True,
+        ),
+        desc.IntParam(
+            name="bondingKernelSize",
+            label="Bonding Kernel Size",
+            description="Kernel size for morphological processing applied for masks bonding.",
+            value=11,
+            enabled=lambda node: node.enableBonding.value,
+        ),
         desc.BoolParam(
             name="maskInvert",
             label="Invert Masks",
@@ -283,7 +296,7 @@ def resolvePaths(self, node):
 
     def processChunk(self, chunk):
         import json
-        from segmentationRDS import image
+        from segmentationRDS import image, sam3Utils
         from sam3.model_builder import build_sam3_video_predictor
         import numpy as np
         import torch
@@ -462,8 +475,16 @@ def processChunk(self, chunk):
                 if len(masks.keys()) > 0:
                     colorPalette.generate_palette(max(masks.keys()) + 1)
                 cryptoName = "object" if prompt == "" else prompt
+
+                if masks.keys():
+                    masks_stack = np.stack(list(masks.values()), axis=0)
+                    bool_mask = np.sum(masks_stack, axis=0) > 0
+                    if len(masks.values()) > 1 and chunk.node.enableBonding.value:
+                        ks = chunk.node.bondingKernelSize.value
+                        bool_mask = sam3Utils.bond_masks(list(masks.values()), ks, ks, ks) > 0
+                    maskImage[bool_mask] = [255, 255, 255]
+
                 for key, mask in masks.items():
-                    maskImage[mask] = [255, 255, 255]
                     color = colorPalette.at(int(key)) if colorPalette.at(int(key)) is not None else [255, 255, 255]
                     colorMaskImage[mask] = [x/255.0 for x in color]
                     if chunk.node.outputCryptomatte.value:
@@ -472,9 +493,15 @@ def processChunk(self, chunk):
                         manifest[obj_name] = hex_val
                         crypto_id[mask] = f32_hash
                         crypto_cov[mask] = 1.0
+
                 if frameId in outputs_per_frame_bwd.keys():
-                    for key, mask in outputs_per_frame_bwd[frameId].items():
-                        maskImage[mask] = [255, 255, 255]
+                    if outputs_per_frame_bwd[frameId].keys():
+                        masks_stack = np.stack(list(outputs_per_frame_bwd[frameId].values()), axis=0)
+                        bool_mask = np.sum(masks_stack, axis=0) > 0
+                        if len(outputs_per_frame_bwd[frameId].values()) > 1 and chunk.node.enableBonding.value:
+                            ks = chunk.node.bondingKernelSize.value
+                            bool_mask = sam3Utils.bond_masks(list(outputs_per_frame_bwd[frameId].values()), ks, ks, ks) > 0
+                        maskImage[bool_mask] = [255, 255, 255]
 
                 if chunk.node.outputCryptomatte.value:
                     spec = oiio.ImageSpec(img.shape[1], img.shape[0], 7, oiio.FLOAT)

meshroom/sam3/VideoSegmentationSam3Boxes.py

@@ -1,4 +1,4 @@
-__version__ = "4.0"
+__version__ = "4.1"
 
 import os
 from pathlib import Path
@@ -82,7 +82,7 @@ class VideoSegmentationSam3Boxes(desc.Node):
 3. Each chunk is split into tiles (if tiling enabled).
 4. Cropped image sequences are fed to the SAM3 video predictor with a text prompt on the first frame;
    masks are propagated across all frames.
-5. Tile masks are resized and composited into full-resolution masks using a **union** operation.
+5. Tile masks are resized and composited into full-resolution masks using a **union** operation and optional bonding.
 6. *(Tiling only)* Tiles with IoU below `minIoU` are replaced by the coarse mask crop.
 7. Masks are saved with optional inversion and bounding box metadata in file headers.
 
@@ -185,6 +185,19 @@ class VideoSegmentationSam3Boxes(desc.Node):
             value=0.5,
             enabled=lambda node: node.enableTiling.value,
         ),
+        desc.BoolParam(
+            name="enableBonding",
+            label="Enable Masks Bonding",
+            description="Enable bonding where instances overlap.",
+            value=True,
+        ),
+        desc.IntParam(
+            name="bondingKernelSize",
+            label="Bonding Kernel Size",
+            description="Kernel size for morphological processing applied for masks bonding.",
+            value=11,
+            enabled=lambda node: node.enableBonding.value,
+        ),
         desc.File(
             name="segmentationModelPath",
             label="Segmentation Model",
@@ -471,12 +484,20 @@ def processChunk(self, chunk):
                                 fine_mask = np.zeros_like(tgt)
                                 frameId = frame_idx - chunk_tile.start_frame
                                 logger.debug(f"frame: {frame_idx}; tile: {box}; items number: {len(outputs_per_frame_visu[frameId].keys())}")
+                                masks = []
                                 for key, maskBoxProb in outputs_per_frame_visu[frameId].items():
                                     mask = maskBoxProb["mask"]
                                     buf_in = oiio.ImageBuf(mask.astype('float32'))
                                     buf_out = oiio.ImageBufAlgo.resample(buf_in, roi=oiio.ROI(0, box_w, 0, box_h))
                                     mask = buf_out.get_pixels().reshape(box_h, box_w, 1)
-                                    bool_mask = mask.squeeze() > 0
+                                    masks.append(mask.squeeze())
+
+                                if masks:
+                                    masks_stack = np.stack(masks, axis=0)
+                                    bool_mask = np.sum(masks_stack, axis=0) > 0
+                                    if len(masks) > 1 and chunk.node.enableBonding.value:
+                                        ks = chunk.node.bondingKernelSize.value
+                                        bool_mask = sam3Utils.bond_masks(masks, ks, ks, ks) > 0
                                     fine_mask[bool_mask] = [255, 255, 255] if maskNbChannel == 3 else [255]
 
                                 if chunk.node.enableTiling.value:

meshroom/sam3/VideoSegmentationSam3Text.py

@@ -61,6 +61,19 @@ class VideoSegmentationSam3Text(desc.Node):
             value=16,
             enabled=lambda node: node.timeSlicing.value,
         ),
+        desc.BoolParam(
+            name="enableBonding",
+            label="Enable Masks Bonding",
+            description="Enable bonding where instances overlap.",
+            value=True,
+        ),
+        desc.IntParam(
+            name="bondingKernelSize",
+            label="Bonding Kernel Size",
+            description="Kernel size for morphological processing applied for masks bonding.",
+            value=11,
+            enabled=lambda node: node.enableBonding.value,
+        ),
         desc.BoolParam(
             name="maskInvert",
             label="Invert Masks",
@@ -360,9 +373,10 @@ def processChunk(self, chunk):
                             crypto_cov_fwd = np.zeros((img.shape[0], img.shape[1]), dtype=np.float32)
                             manifest_fwd = {}
                         boxes[textPrompt]["forward"][firstFrameId + frameId] = {}
+                        masks = []
                         for key, maskBoxProb in outputs_per_frame_visu[frameId].items():
                             mask = maskBoxProb["mask"]
-                            mask_images[frameId][mask] = [(int(key) + 1) * 255, 255, 255]
+                            masks.append(mask.squeeze())
                             color = colorPalette.at(int(key)) if colorPalette.at(int(key)) is not None else [255, 255, 255]
                             colorMaskImageFwd[mask] = [x/255.0 for x in color]
 
@@ -379,6 +393,14 @@ def processChunk(self, chunk):
                             bbox_str = str(x1) + ";" + str(y1) + ";" + str(x2) + ";" + str(y2)
                             metadata_boxes[frameId][textPrompt]["forward"]["fwd_" + textPrompt + "_" + str(key)] = bbox_str
 
+                        if masks:
+                            masks_stack = np.stack(masks, axis=0)
+                            mask_global = np.expand_dims(np.sum(masks_stack, axis=0), axis=-1)
+                            if len(masks) > 1 and chunk.node.enableBonding.value:
+                                ks = chunk.node.bondingKernelSize.value
+                                mask_global = np.expand_dims(sam3Utils.bond_masks(masks, ks, ks, ks), axis=-1)
+                            mask_images[frameId] = mask_global
+
                         if chunk.node.outputColorMasks.value:
                             if chunk.node.keepFilename.value:
                                 outputFileColorMask = os.path.join(chunk.node.output.value, "colorMask_" + textPrompt + "_fwd_" + str(Path(chunk_image_paths[frameId][0]).stem) + ".png")
@@ -410,9 +432,10 @@ def processChunk(self, chunk):
                                 crypto_cov_bwd = np.zeros((img.shape[0], img.shape[1]), dtype=np.float32)
                                 manifest_bwd = {}
                             boxes[textPrompt]["backward"][firstFrameId + frameId] = {}
+                            masks = []
                             for key, maskBoxProb in outputs_per_frame_visu[frameId].items():
                                 mask = maskBoxProb["mask"]
-                                mask_images[frameId][mask] = [(int(key) + 1) * 255, 255, 255]
+                                masks.append(mask.squeeze())
                                 color = colorPalette.at(int(key)) if colorPalette.at(int(key)) is not None else [255, 255, 255]
                                 colorMaskImageBwd[mask] = [x/255.0 for x in color]
                                 if chunk.node.outputCryptomatte.value:
@@ -427,6 +450,14 @@ def processChunk(self, chunk):
                                 bbox_str = str(x1)+";"+str(y1)+";"+str(x2)+";"+str(y2)
                                 metadata_boxes[frameId][textPrompt]["backward"]["bwd_"+textPrompt+"_"+str(key)] = bbox_str
 
+                            if masks:
+                                masks_stack = np.stack(masks, axis=0)
+                                mask_global = np.expand_dims(np.sum(masks_stack, axis=0), axis=-1)
+                                if len(masks) > 1 and chunk.node.enableBonding.value:
+                                    ks = chunk.node.bondingKernelSize.value
+                                    mask_global = np.expand_dims(sam3Utils.bond_masks(masks, ks, ks, ks), axis=-1)
+                                mask_images[frameId] = mask_global
+
                             if chunk.node.outputColorMasks.value:
                                 if chunk.node.keepFilename.value:
                                     outputFileColorMask = os.path.join(chunk.node.output.value, "colorMask_" + textPrompt + "_bwd_" + str(Path(chunk_image_paths[frameId][0]).stem) + ".png")

segmentationRDS/sam3Utils.py

@@ -132,3 +132,31 @@ def mapIds(outputs_per_frame_detected, outputs_per_frame_propagated, logger):
                 else:
                     mapping[idSrc] = idSrc
     return mapping
+
+def bond_masks(masks, dilate_kernel_size: int = 11, conflict_kernel_size: int = 11, close_kernel_size: int = 11):
+    import cv2
+
+    # Dilate each individual mask
+    kernel_dilate = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilate_kernel_size, dilate_kernel_size))
+    dilated_masks = [cv2.dilate(m.astype(np.uint8), kernel_dilate) for m in masks]
+    dilated_stack = np.stack(dilated_masks, axis=0)
+
+    # Detect overlapping regions to generate the activation mask
+    overlap_map = np.sum(dilated_stack, axis=0)
+    conflict_zone = (overlap_map >= 2).astype(np.uint8) * 255
+    kernel_conflict = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (conflict_kernel_size, conflict_kernel_size))
+    action_mask = cv2.dilate(conflict_zone, kernel_conflict)
+
+    # Merge individual masks into a single global mask
+    masks_stack = np.stack(masks, axis=0)
+    global_mask_raw = (np.sum(masks_stack, axis=0) > 0).astype(np.uint8) * 255
+
+    # Apply morphological closing to fill holes in the global mask
+    kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (close_kernel_size, close_kernel_size))
+    global_mask_closed = cv2.morphologyEx(global_mask_raw, cv2.MORPH_CLOSE, kernel_close)
+
+    # Apply the closed mask only within detected conflict regions
+    bonded_mask = np.where(action_mask > 0, global_mask_closed, global_mask_raw)
+
+    return bonded_mask
+