From f035c0f890a3ed7d2202ea76d2d00ea34cd81224 Mon Sep 17 00:00:00 2001 From: Pablerdo Date: Tue, 25 Feb 2025 17:25:47 -0500 Subject: [PATCH] fix for dimension and tendor dimensions --- .gitignore | 3 +- nodes/curve_nodes.py | 87 +++++++++++++++++++------ nodes/mask_nodes.py | 151 ++++++++++++++++++++++++++++++++++--------- requirements.txt | 10 ++- 4 files changed, 197 insertions(+), 54 deletions(-) diff --git a/.gitignore b/.gitignore index f37158a..eefb916 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ test/ __pycache__ -*.gif \ No newline at end of file +*.gif +nodes/loader_nodes.py \ No newline at end of file diff --git a/nodes/curve_nodes.py b/nodes/curve_nodes.py index b02c307..dea970d 100644 --- a/nodes/curve_nodes.py +++ b/nodes/curve_nodes.py @@ -47,27 +47,37 @@ class MultiCutAndDragOnPath: } def multi_cutanddrag(self, image, coordinate_paths, masks, frame_width, frame_height, inpaint, rotation=False, bg_image=None, degrees=[0.0]): + # Handle case where masks is a tuple (common in ComfyUI node system) + masks_tensor = masks + if isinstance(masks, tuple): + masks_tensor = masks[0] # Extract the mask tensor from the tuple + if not rotation: return self._translate(image, coordinate_paths, masks, frame_width, frame_height, inpaint, bg_image) else: # Verify that degrees array matches number of masks - if len(degrees) != masks.shape[0]: - raise ValueError(f"Number of rotation degrees ({len(degrees)}) must match number of masks ({masks.shape[0]})") + if len(degrees) != masks_tensor.shape[0]: + raise ValueError(f"Number of rotation degrees ({len(degrees)}) must match number of masks ({masks_tensor.shape[0]})") return self._translate_and_rotate(image, coordinate_paths, masks, frame_width, frame_height, inpaint, degrees, bg_image) def _translate(self, image, coordinate_paths, masks, frame_width, frame_height, inpaint, bg_image=None): # Parse coordinate paths as array of arrays paths_list = json.loads(coordinate_paths) - if len(paths_list) != masks.shape[0]: - raise ValueError(f"Number of coordinate paths ({len(paths_list)}) must match number of masks ({masks.shape[0]})") + # Handle case where masks is a tuple + masks_tensor = masks + if isinstance(masks, tuple): + masks_tensor = masks[0] + + if len(paths_list) != masks_tensor.shape[0]: + raise ValueError(f"Number of coordinate paths ({len(paths_list)}) must match number of masks ({masks_tensor.shape[0]})") batch_size = len(paths_list[0]) # Number of frames to generate images_list = [] masks_list = [] # Convert input image to PIL - input_image = tensor2pil(image)[0] + input_image = tensor2pil(image[0])[0] # Create inpainted background once if needed if bg_image is None: @@ -79,9 +89,24 @@ class MultiCutAndDragOnPath: draw = ImageDraw.Draw(combined_mask) border = 5 - for mask_idx in range(masks.shape[0]): - mask_pil = tensor2pil(masks[mask_idx:mask_idx+1])[0] - mask_array = np.array(mask_pil) + for mask_idx in range(masks_tensor.shape[0]): + # Get mask and ensure it's properly shaped for PIL conversion + mask_tensor = masks_tensor[mask_idx] + + # Handle potentially problematic dimensions + if len(mask_tensor.shape) == 3 and mask_tensor.shape[0] == 1: + mask_tensor = mask_tensor.squeeze(0) # Remove singleton dimension + + # Convert to numpy array directly instead of using tensor2pil + mask_array = mask_tensor.cpu().numpy() + + # Ensure mask_array is 2D + if len(mask_array.shape) == 3: + if mask_array.shape[0] == 1: + mask_array = mask_array[0] # Take first channel + else: + mask_array = mask_array.mean(axis=0) # Average channels + y_indices, x_indices = np.where(mask_array > 0) if len(x_indices) > 0 and len(y_indices) > 0: x_min, x_max = x_indices.min(), x_indices.max() @@ -100,9 +125,28 @@ class MultiCutAndDragOnPath: # Cut out each masked region and store info cut_regions = [] - for mask_idx in range(masks.shape[0]): - mask_pil = tensor2pil(masks[mask_idx:mask_idx+1])[0] - mask_array = np.array(mask_pil) + for mask_idx in range(masks_tensor.shape[0]): + # Get mask and ensure it's properly shaped + mask_tensor = masks_tensor[mask_idx] + + # Handle potentially problematic dimensions + if len(mask_tensor.shape) == 3 and mask_tensor.shape[0] == 1: + mask_tensor = mask_tensor.squeeze(0) # Remove singleton dimension + + # Convert to numpy array directly + mask_array = mask_tensor.cpu().numpy() + + # Ensure mask_array is 2D + if len(mask_array.shape) == 3: + if mask_array.shape[0] == 1: + mask_array = mask_array[0] # Take first channel + else: + mask_array = mask_array.mean(axis=0) # Average channels + + # Create PIL mask directly from the array + mask_pil = Image.fromarray((mask_array * 255).astype(np.uint8)) + + # Rest of the processing remains the same y_indices, x_indices = np.where(mask_array > 0) if len(x_indices) > 0 and len(y_indices) > 0: @@ -164,12 +208,17 @@ class MultiCutAndDragOnPath: except ValueError: raise ValueError("Degrees must be a valid number or JSON array of numbers") - if len(paths_list) != masks.shape[0]: - raise ValueError(f"Number of coordinate paths ({len(paths_list)}) must match number of masks ({masks.shape[0]})") + # Handle case where masks is a tuple + masks_tensor = masks + if isinstance(masks, tuple): + masks_tensor = masks[0] + + if len(paths_list) != masks_tensor.shape[0]: + raise ValueError(f"Number of coordinate paths ({len(paths_list)}) must match number of masks ({masks_tensor.shape[0]})") # Verify that degrees array matches number of masks - if len(degrees_list) != masks.shape[0]: - raise ValueError(f"Number of rotation degrees ({len(degrees_list)}) must match number of masks ({masks.shape[0]})") + if len(degrees_list) != masks_tensor.shape[0]: + raise ValueError(f"Number of rotation degrees ({len(degrees_list)}) must match number of masks ({masks_tensor.shape[0]})") batch_size = len(paths_list[0]) # Number of frames to generate images_list = [] @@ -188,8 +237,8 @@ class MultiCutAndDragOnPath: draw = ImageDraw.Draw(combined_mask) border = 5 - for mask_idx in range(masks.shape[0]): - mask_pil = tensor2pil(masks[mask_idx:mask_idx+1])[0] + for mask_idx in range(masks_tensor.shape[0]): + mask_pil = tensor2pil(masks_tensor[mask_idx:mask_idx+1])[0] mask_array = np.array(mask_pil) y_indices, x_indices = np.where(mask_array > 0) if len(x_indices) > 0 and len(y_indices) > 0: @@ -209,8 +258,8 @@ class MultiCutAndDragOnPath: # Cut out each masked region and store info cut_regions = [] - for mask_idx in range(masks.shape[0]): - mask_pil = tensor2pil(masks[mask_idx:mask_idx+1])[0] + for mask_idx in range(masks_tensor.shape[0]): + mask_pil = tensor2pil(masks_tensor[mask_idx:mask_idx+1])[0] mask_array = np.array(mask_pil) y_indices, x_indices = np.where(mask_array > 0) diff --git a/nodes/mask_nodes.py b/nodes/mask_nodes.py index 7951ffe..3cecf1c 100644 --- a/nodes/mask_nodes.py +++ b/nodes/mask_nodes.py @@ -21,6 +21,8 @@ class BatchImageToMask: "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}), + "save_to_file": ("BOOLEAN", {"default": False}), + "output_path": ("STRING", {"default": "masks_output.png"}) }} CATEGORY = "PSNodes/masking" @@ -29,30 +31,57 @@ class BatchImageToMask: FUNCTION = "batch_convert_to_mask" DESCRIPTION = "Converts RGB images to binary masks using grayscale conversion and thresholding, with optional morphological dilation" - def batch_convert_to_mask(self, images, threshold, dilation_amount): - # Get shape - handle both [B, H, W] and [B, C, H, W] formats - if len(images.shape) == 3: - B, H, W = images.shape # Already in [B, H, W] format - # Create a tensor to hold the masks - masks_tensor = torch.zeros((B, H, W), device=images.device) - - # Convert each image to a binary mask - for i in range(B): - mask = (images[i] > threshold).float() # Binarize using threshold - masks_tensor[i] = mask - else: - # Handle [B, C, H, W] format (standard ComfyUI IMAGE) + def batch_convert_to_mask(self, images, threshold, dilation_amount, save_to_file=False, output_path="masks_output.png"): + # Handle case where images is a tuple (common in ComfyUI node system) + if isinstance(images, tuple): + images = images[0] + + # Print the shape to diagnose + print(f"Input images shape: {images.shape}") + + # Check the tensor shape and transpose if needed + if len(images.shape) == 4: # [B, C, H, W] B, C, H, W = images.shape - masks_tensor = torch.zeros((B, H, W), device=images.device) - - # Convert each RGB image to a binary mask + if C == 3: + # Standard RGB format + pass + elif W == 3 and C > 3: + # Dimensions are swapped [B, H, W, C] + print("Dimensions detected as [B, H, W, C], correcting to [B, C, H, W]") + # Use correct permutation: from [B, H, W, C] to [B, C, H, W] + images = images.permute(0, 3, 1, 2) + B, C, H, W = images.shape + elif len(images.shape) == 3: # Could be [B, H, W] or [H, W, C] + if images.shape[2] == 3: # [H, W, C] format + print("Transposing image tensor from [H, W, C] to [1, C, H, W]") + # Correct permutation: from [H, W, C] to [1, C, H, W] + images = images.permute(2, 0, 1).unsqueeze(0) + B, C, H, W = images.shape + else: + # Assume [B, H, W] format + B, H, W = images.shape + C = 1 + + print(f"After shape analysis: B={B}, C={C if 'C' in locals() else 1}, H={H}, W={W}") + + # Create a properly shaped tensor for masks + masks = torch.zeros((B, C, H, W), device=images.device) + + # Convert each image to a binary mask + if len(images.shape) == 4: # [B, C, H, W] for i in range(B): - img = images[i] # Shape: [C, H, W] - # Convert to grayscale by taking mean across channels - gray = img.mean(dim=0) # Shape: [H, W] - mask = (gray > threshold).float() # Binarize using threshold - masks_tensor[i] = mask - + # Convert RGB to grayscale properly + if C == 3: + # Standard RGB format + gray = images[i].mean(dim=0) + else: + # Single channel + gray = images[i][0] + masks[i] = (gray > threshold).float() + else: # [B, H, W] + for i in range(B): + masks[i] = (images[i] > threshold).float() + # Apply morphological dilation if requested if dilation_amount > 0: c = 1 # non-tapered corners @@ -61,15 +90,58 @@ class BatchImageToMask: [c, 1, c]]) dilated_masks = [] - for m in masks_tensor: + for m in masks: output = m.cpu().numpy().astype(np.float32) for _ in range(dilation_amount): output = scipy.ndimage.grey_dilation(output, footprint=kernel) dilated_masks.append(torch.from_numpy(output)) - masks_tensor = torch.stack(dilated_masks).to(images.device) + masks = torch.stack(dilated_masks).to(images.device) - return (masks_tensor,) + # # Save masks to file if requested + # if save_to_file: + # # Create a large image to hold all masks side by side or in a grid + # if B == 1: + # # Just one mask, save it directly + # mask_pil_result = tensor2pil(masks[0]) + # # Handle case where tensor2pil returns a list + # if isinstance(mask_pil_result, list): + # mask_pil = mask_pil_result[0] + # print(f"PIL mask dimensions: {mask_pil.size}") + # else: + # mask_pil = mask_pil_result + # print(f"PIL mask dimensions: {mask_pil.size}") + + # mask_pil = mask_pil.convert("L") # Convert to grayscale + # mask_pil.save(output_path) + # else: + # # Multiple masks, arrange in a grid + # rows = int(np.ceil(np.sqrt(B))) + # cols = int(np.ceil(B / rows)) + # grid_image = Image.new("L", (cols * W, rows * H), 0) + + # for i in range(B): + # mask_pil_result = tensor2pil(masks[i]) + # # Handle case where tensor2pil returns a list + # if isinstance(mask_pil_result, list): + # mask_pil = mask_pil_result[0] + # print(f"PIL mask dimensions: {mask_pil.size}") + # else: + # mask_pil = mask_pil_result + # print(f"PIL mask dimensions: {mask_pil.size}") + + # mask_pil = mask_pil.convert("L") # Convert to grayscale + # row = i // cols + # col = i % cols + # grid_image.paste(mask_pil, (col * W, row * H)) + + # grid_image.save(output_path) + # print(f"Saved {B} masks to {output_path}") + + print(f"Mask tensor shape: {masks.shape}") + print(f"Mask value range: {masks.min()} to {masks.max()}") + + return (masks,) class MapTrajectoriesToSegmentedMasks: @classmethod @@ -96,21 +168,38 @@ class MapTrajectoriesToSegmentedMasks: 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 + # Get dimensions, handling different tensor formats + if len(masks.shape) == 4: # [B, C, H, W] + B, C, H, W = masks.shape + elif len(masks.shape) == 3: # [B, H, W] + B, H, W = masks.shape + C = 1 + else: + raise ValueError(f"Unexpected mask shape: {masks.shape}") # Process each mask and trajectory translated_paths = [] for i in range(B): - # Get the mask + # Get the mask and convert to numpy mask = masks[i] - - # Calculate centroid mask_np = mask.cpu().numpy() + + # Handle masks with different dimensions + if len(mask_np.shape) > 2: # Multi-channel mask [C, H, W] + # Average across channels if needed + if mask_np.shape[0] > 1: + mask_np = mask_np.mean(axis=0) + else: + mask_np = mask_np[0] # Take first channel + + # Calculate centroid correctly on 2D mask if mask_np.sum() > 0: - y_indices, x_indices = np.where(mask_np > 0.5) - centroid_x = x_indices.mean() + # This works on 2D arrays + indices = np.where(mask_np > 0.5) + y_indices, x_indices = indices[0], indices[1] centroid_y = y_indices.mean() + centroid_x = x_indices.mean() else: # Default to center if no mask points centroid_x = W / 2 diff --git a/requirements.txt b/requirements.txt index 5bc18ca..a8c5cee 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,11 @@ -pillow>=10.3.0 -scipy +torch>=2.0.0 +torchvision>=0.15.0 +numpy>=1.20.0 +Pillow>=9.0.0 +scipy>=1.8.0 color-matcher matplotlib huggingface_hub mss -opencv-python \ No newline at end of file +opencv-python +comfy \ No newline at end of file