nodes/Pablerdo-ComfyUI-MultiCutAndDrag
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix for dimension and tendor dimensions

Browse Source
This commit is contained in:
Pablerdo
2025-02-25 17:25:47 -05:00
parent 01a5bd2a09
commit f035c0f890
4 changed files with 197 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2 -1
View File
@@ -1,3 +1,4 @@
test/
__pycache__
*.gif
*.gif
nodes/loader_nodes.py
nodes/curve_nodes.py
+68 -19
View File
@@ -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)
nodes/mask_nodes.py
+120 -31
View File
@@ -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
requirements.txt
+7 -3
View File
@@ -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
opencv-python
comfy