diff --git a/meshroom/sam3/ImageSegmentationSam3.py b/meshroom/sam3/ImageSegmentationSam3.py index adec7b2..3530856 100644 --- a/meshroom/sam3/ImageSegmentationSam3.py +++ b/meshroom/sam3/ImageSegmentationSam3.py @@ -51,6 +51,20 @@ 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, + range=(1, 255, 2), + enabled=lambda node: node.enableBonding.value, + ), desc.BoolParam( name="maskInvert", label="Invert Masks", @@ -238,11 +252,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 +369,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') diff --git a/meshroom/sam3/VideoSegmentationSam3.py b/meshroom/sam3/VideoSegmentationSam3.py index f050017..7d32909 100644 --- a/meshroom/sam3/VideoSegmentationSam3.py +++ b/meshroom/sam3/VideoSegmentationSam3.py @@ -62,6 +62,20 @@ 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, + range=(1, 255, 2), + enabled=lambda node: node.enableBonding.value, + ), desc.BoolParam( name="maskInvert", label="Invert Masks", @@ -283,7 +297,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 +476,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 +494,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) diff --git a/meshroom/sam3/VideoSegmentationSam3Boxes.py b/meshroom/sam3/VideoSegmentationSam3Boxes.py index 70e6f57..e00e983 100644 --- a/meshroom/sam3/VideoSegmentationSam3Boxes.py +++ b/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,20 @@ 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, + range=(1, 255, 2), + enabled=lambda node: node.enableBonding.value, + ), desc.File( name="segmentationModelPath", label="Segmentation Model", @@ -471,12 +485,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: diff --git a/meshroom/sam3/VideoSegmentationSam3Text.py b/meshroom/sam3/VideoSegmentationSam3Text.py index 9024df6..3ef9c9a 100644 --- a/meshroom/sam3/VideoSegmentationSam3Text.py +++ b/meshroom/sam3/VideoSegmentationSam3Text.py @@ -61,6 +61,20 @@ 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, + range=(1, 255, 2), + enabled=lambda node: node.enableBonding.value, + ), desc.BoolParam( name="maskInvert", label="Invert Masks", @@ -360,9 +374,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 +394,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 +433,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 +451,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") diff --git a/segmentationRDS/sam3Utils.py b/segmentationRDS/sam3Utils.py index fb99067..6f77568 100644 --- a/segmentationRDS/sam3Utils.py +++ b/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(np.round(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 + \ No newline at end of file