From 083201ca2d0c85f4f1ea612c4dadf092e7ae3e83 Mon Sep 17 00:00:00 2001 From: Pablerdo Date: Mon, 24 Feb 2025 23:50:24 -0500 Subject: [PATCH] map trajectories node --- .gitignore | 3 ++ __init__.py | 3 ++ nodes/curve_nodes.py | 2 - nodes/image_nodes.py | 2 +- nodes/mask_nodes.py | 96 +++++++++++++++++++++++++++++++++++--------- 5 files changed, 84 insertions(+), 22 deletions(-) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f37158a --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +test/ +__pycache__ +*.gif \ No newline at end of file diff --git a/__init__.py b/__init__.py index 3f588ea..85af893 100644 --- a/__init__.py +++ b/__init__.py @@ -2,12 +2,15 @@ from .nodes.curve_nodes import MultiCutAndDragOnPath from .nodes.mask_nodes import BatchImageToMask from .nodes.image_nodes import LoadImagesFromBase64Array from .nodes.image_nodes import LoadImageFromBase64 +from .nodes.mask_nodes import MapTrajectoriesToSegmentedMasks NODE_CONFIG = { "LoadImageFromBase64": {"class": LoadImageFromBase64, "name": "LoadImageFromBase64"}, "LoadImagesFromBase64Array": {"class": LoadImagesFromBase64Array, "name": "LoadImagesFromBase64Array"}, "MultiCutAndDragOnPath": {"class": MultiCutAndDragOnPath, "name": "MultiCutAndDragOnPath"}, "BatchImageToMask": {"class": BatchImageToMask, "name": "BatchImageToMask"}, + "MapTrajectoriesToSegmentedMasks": {"class": MapTrajectoriesToSegmentedMasks, "name": "MapTrajectoriesToSegmentedMasks"} + } def generate_node_mappings(node_config): diff --git a/nodes/curve_nodes.py b/nodes/curve_nodes.py index 84065d0..b02c307 100644 --- a/nodes/curve_nodes.py +++ b/nodes/curve_nodes.py @@ -10,11 +10,9 @@ import sys sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from utility.utility import pil2tensor, tensor2pil -import folder_paths import io import base64 -from comfy.utils import common_upscale class MultiCutAndDragOnPath: RETURN_TYPES = ("IMAGE", "MASK",) diff --git a/nodes/image_nodes.py b/nodes/image_nodes.py index 14a75a2..36a71f8 100644 --- a/nodes/image_nodes.py +++ b/nodes/image_nodes.py @@ -3,7 +3,7 @@ import os import cv2 from PIL import Image, ImageOps -import folder_paths +# import folder_paths import torch import numpy as np from io import BytesIO diff --git a/nodes/mask_nodes.py b/nodes/mask_nodes.py index 8835407..ddd81b4 100644 --- a/nodes/mask_nodes.py +++ b/nodes/mask_nodes.py @@ -6,14 +6,10 @@ import scipy.ndimage import numpy as np from contextlib import nullcontext import os - +import json import model_management -from comfy.utils import ProgressBar -from comfy.utils import common_upscale from nodes import MAX_RESOLUTION -import folder_paths - from ..utility.utility import tensor2pil, pil2tensor script_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -23,6 +19,7 @@ class BatchImageToMask: def INPUT_TYPES(s): return {"required": { "images": ("IMAGE",), + "threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), "dilation_amount": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), }} @@ -30,22 +27,20 @@ class BatchImageToMask: RETURN_TYPES = ("MASK",) RETURN_NAMES = ("MASKS",) FUNCTION = "batch_convert_to_mask" - DESCRIPTION = "Converts transparent PNG images to binary masks with optional dilation" + DESCRIPTION = "Converts RGB images to binary masks using grayscale conversion and thresholding, with optional morphological dilation" - def batch_convert_to_mask(self, images, dilation_amount): + def batch_convert_to_mask(self, images, threshold, dilation_amount): B, H, W, C = images.shape - masks = [] - - # Convert images to masks based on alpha channel + masks_tensor = torch.zeros((B, H, W), device=images.device) + + # Convert each RGB image to a binary mask for i in range(B): - # Extract alpha channel (assuming RGBA where A is the last channel) - mask = images[i, :, :, -1] # Get alpha channel - mask = (mask > 0.5).float() # Binarize the mask - masks.append(mask) - - mask_tensor = torch.stack(masks) - - # Apply morphological dilation if requested + img = images[i] + gray = img.mean(dim=-1) # Simple RGB to grayscale conversion + mask = (gray > threshold).float() # Binarize using threshold + masks_tensor[i] = mask + + # Apply morphological dilation if requested if dilation_amount > 0: c = 1 # non-tapered corners kernel = np.array([[c, 1, c], @@ -61,4 +56,67 @@ class BatchImageToMask: mask_tensor = torch.stack(dilated_masks) - return (mask_tensor,) \ No newline at end of file + return (masks_tensor,) + +class MapTrajectoriesToSegmentedMasks: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "masks": ("MASK",), + "trajectories": ("STRING", {"forceInput": True}), + }} + + CATEGORY = "PSNodes/masking" + RETURN_TYPES = ("MASK", "STRING") + RETURN_NAMES = ("MASKS", "TRANSLATED_TRAJECTORIES") + FUNCTION = "map_trajectories" + DESCRIPTION = "Maps trajectories to masks by centering the start point of trajectories on mask centroids" + + def map_trajectories(self, masks, trajectories): + # Parse trajectories JSON + paths_list = json.loads(trajectories) + + if len(paths_list) != masks.shape[0]: + raise ValueError(f"Number of trajectories ({len(paths_list)}) must match number of masks ({masks.shape[0]})") + + B, H, W = masks.shape + + # Process each mask and trajectory + translated_paths = [] + + for i in range(B): + # Get the mask + mask = masks[i] + + # Calculate centroid + mask_np = mask.cpu().numpy() + if mask_np.sum() > 0: + y_indices, x_indices = np.where(mask_np > 0.5) + centroid_x = x_indices.mean() + centroid_y = y_indices.mean() + else: + # Default to center if no mask points + centroid_x = W / 2 + centroid_y = H / 2 + + # Get corresponding trajectory and translate it + path = paths_list[i] + + # Calculate the offset from the first point to the centroid + if len(path) > 0: + first_point = path[0] + offset_x = centroid_x - first_point["x"] + offset_y = centroid_y - first_point["y"] + + # Apply the offset to all points in the path + translated_path = [{"x": point["x"] + offset_x, "y": point["y"] + offset_y} for point in path] + translated_paths.append(translated_path) + else: + # Empty path case + translated_paths.append([]) + + # Convert translated paths back to JSON + translated_trajectories_json = json.dumps(translated_paths) + + return (masks, translated_trajectories_json) +