xinyun/ComfyUI-KJNodes
1
0
Fork 0
mirror of https://git.datalinker.icu/kijai/ComfyUI-KJNodes.git synced 2025-12-15 07:44:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update mask_nodes.py

Browse Source
This commit is contained in:
kijai 2025-07-03 20:45:10 +03:00
parent 1f653a8d96
commit 40d9d84b68
1 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
nodes/mask_nodes.py
View File

@ -1304,7 +1304,7 @@ class SeparateMasks:
"mask": ("MASK", ), "mask": ("MASK", ),
"size_threshold_width" : ("INT", {"default": 256, "min": 0.0, "max": 4096, "step": 1}), "size_threshold_width" : ("INT", {"default": 256, "min": 0.0, "max": 4096, "step": 1}),
"size_threshold_height" : ("INT", {"default": 256, "min": 0.0, "max": 4096, "step": 1}), "size_threshold_height" : ("INT", {"default": 256, "min": 0.0, "max": 4096, "step": 1}),
"mode": (["convex_polygons", "area"],), "mode": (["convex_polygons", "area", "box"],),
"max_poly_points": ("INT", {"default": 8, "min": 3, "max": 32, "step": 1}), "max_poly_points": ("INT", {"default": 8, "min": 3, "max": 32, "step": 1}),
}, },
@ -1399,12 +1399,18 @@ class SeparateMasks:
print(f"Component {component}: width={width}, height={height}, x_pos={centroid_x}") print(f"Component {component}: width={width}, height={height}, x_pos={centroid_x}")
if width >= size_threshold_width and height >= size_threshold_height: if width >= size_threshold_width and height >= size_threshold_height:
if mode != "area": if mode == "convex_polygons":
polygon = self.get_mask_polygon(component_mask_np, max_poly_points) polygon = self.get_mask_polygon(component_mask_np, max_poly_points)
if polygon is not None: if polygon is not None:
poly_mask = self.polygon_to_mask(polygon, (H, W)) poly_mask = self.polygon_to_mask(polygon, (H, W))
poly_mask = torch.tensor(poly_mask, device=mask.device) poly_mask = torch.tensor(poly_mask, device=mask.device)
separated.append((centroid_x, poly_mask)) separated.append((centroid_x, poly_mask))
elif mode == "box":
# Create bounding box mask
box_mask = np.zeros((H, W), dtype=np.uint8)
box_mask[y_min:y_max+1, x_min:x_max+1] = 1
box_mask = torch.tensor(box_mask, device=mask.device)
separated.append((centroid_x, box_mask))
else: else:
area_mask = torch.tensor(component_mask_np, device=mask.device) area_mask = torch.tensor(component_mask_np, device=mask.device)
separated.append((centroid_x, area_mask)) separated.append((centroid_x, area_mask))