map trajectories node

.gitignore

@@ -0,0 +1,3 @@
+test/
+__pycache__
+*.gif

__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):

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",)

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

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,)
+        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)
+