Initial commit
This commit is contained in:
TechnicalTod
+187
@@ -0,0 +1,187 @@
|
||||
"""
|
||||
Roto Mask Node - Backend implementation
|
||||
Converts curve data to rasterized masks
|
||||
"""
|
||||
|
||||
import json
|
||||
import numpy as np
|
||||
import torch
|
||||
from typing import Dict, List, Tuple, Optional, Any
|
||||
|
||||
from .core.curve_validator import validate_schema, validate_curve
|
||||
from .mask.generator import generate_single_mask
|
||||
from .image.loader import load_image_from_disk
|
||||
from .image.processor import resize_mask_to_image, apply_mask_overlay
|
||||
from .image.utils import create_empty_mask
|
||||
|
||||
|
||||
class RotoMaskNode:
|
||||
"""
|
||||
ComfyUI node for generating masks from curve data.
|
||||
Accepts image input, curve JSON data, and frame index.
|
||||
Outputs a single-channel mask image.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> Dict[str, Dict[str, Any]]:
|
||||
"""
|
||||
Define input parameters for the node.
|
||||
|
||||
Returns:
|
||||
Dictionary specifying required input types and their constraints
|
||||
"""
|
||||
return {
|
||||
"required": {
|
||||
"curve_json": (
|
||||
"STRING",
|
||||
{
|
||||
"multiline": True,
|
||||
"default": '{"meta":{"version":1,"width":512,"height":512},"shapes":{}}',
|
||||
},
|
||||
),
|
||||
"frame_index": ("INT", {"default": 1, "min": 1, "max": 9999, "step": 1}),
|
||||
},
|
||||
"optional": {
|
||||
"frame_count": ("INT", {"default": 1, "min": 1, "max": 9999, "step": 1}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE", "IMAGE", "STRING")
|
||||
RETURN_NAMES = ("Mask as image", "Mask", "Masked Image", "Loaded Image", "Mask Data")
|
||||
FUNCTION = "generate_mask"
|
||||
CATEGORY = "mask"
|
||||
|
||||
def generate_mask(
|
||||
self, curve_json: str, frame_index: int, frame_count: Optional[int] = None
|
||||
) -> Tuple[torch.Tensor]:
|
||||
"""
|
||||
Main execution function for the node.
|
||||
|
||||
Args:
|
||||
curve_json: JSON string containing curve data
|
||||
frame_index: Starting frame number to render
|
||||
frame_count: Number of frames to render (1 = single frame, >1 = batch sequence)
|
||||
|
||||
Returns:
|
||||
Tuple containing mask tensor as RGB image (B, H, W, 3)
|
||||
"""
|
||||
# Default frame_count to 1 if not provided
|
||||
if frame_count is None:
|
||||
frame_count = 1
|
||||
|
||||
# For batch rendering (frame_count > 1), always start from frame 1
|
||||
# frame_index is only used for single-frame preview mode
|
||||
if frame_count > 1:
|
||||
frame_index = 1
|
||||
|
||||
# Parse curve data
|
||||
try:
|
||||
curve_data = json.loads(curve_json)
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"[RotoMaskNode] Invalid JSON: {e}")
|
||||
# Return empty masks on error (default 512x512)
|
||||
empty_mask = create_empty_mask(frame_count, 512, 512)
|
||||
empty_mask_single = torch.zeros((frame_count, 512, 512), dtype=torch.float32)
|
||||
return (empty_mask, empty_mask_single, empty_mask, empty_mask, curve_json)
|
||||
|
||||
# Validate schema
|
||||
if not validate_schema(curve_data):
|
||||
print("[RotoMaskNode] Invalid curve data schema")
|
||||
# Return empty masks on error (default 512x512)
|
||||
empty_mask = create_empty_mask(frame_count, 512, 512)
|
||||
empty_mask_single = torch.zeros((frame_count, 512, 512), dtype=torch.float32)
|
||||
return (empty_mask, empty_mask_single, empty_mask, empty_mask, curve_json)
|
||||
|
||||
# Get dimensions from curve_json metadata, default to 512x512
|
||||
meta = curve_data.get("meta", {})
|
||||
width = meta.get("width", 512)
|
||||
height = meta.get("height", 512)
|
||||
|
||||
# Get feather_amount from metadata, validate and clamp
|
||||
feather_amount = meta.get("feather_amount", 0.0)
|
||||
if not isinstance(feather_amount, (int, float)) or not np.isfinite(feather_amount):
|
||||
feather_amount = 0.0
|
||||
else:
|
||||
feather_amount = float(np.clip(feather_amount, 0.0, 100.0))
|
||||
|
||||
# Get background_color from metadata, default to black
|
||||
background_color = meta.get("background_color", [0.0, 0.0, 0.0])
|
||||
if not isinstance(background_color, list) or len(background_color) != 3:
|
||||
background_color = [0.0, 0.0, 0.0]
|
||||
else:
|
||||
background_color = [
|
||||
float(np.clip(background_color[0], 0.0, 1.0)),
|
||||
float(np.clip(background_color[1], 0.0, 1.0)),
|
||||
float(np.clip(background_color[2], 0.0, 1.0)),
|
||||
]
|
||||
|
||||
# Get frame_paths from metadata
|
||||
frame_paths = meta.get("frame_paths", [])
|
||||
|
||||
# Generate masks for all frames in the sequence
|
||||
masks = []
|
||||
for i in range(frame_count):
|
||||
current_frame = frame_index + i
|
||||
mask = generate_single_mask(curve_data, current_frame, height, width, feather_amount)
|
||||
masks.append(mask)
|
||||
|
||||
# Stack all masks into batch: list of (H, W) -> (B, H, W)
|
||||
mask_batch = np.stack(masks, axis=0)
|
||||
|
||||
# Output 1: Convert to 3-channel RGB for preview compatibility: (B, H, W) -> (B, H, W, 3)
|
||||
mask_rgb = np.stack([mask_batch, mask_batch, mask_batch], axis=-1)
|
||||
mask_tensor = torch.from_numpy(mask_rgb).float()
|
||||
|
||||
# Output 2: MASK format (single channel, B, H, W)
|
||||
mask_single_channel = torch.from_numpy(mask_batch).float()
|
||||
|
||||
# Outputs 3 and 4: Load images and apply mask overlay
|
||||
masked_images = []
|
||||
raw_images = []
|
||||
|
||||
for i in range(frame_count):
|
||||
current_frame = frame_index + i
|
||||
# Map 1-based frame_index to 0-based array index
|
||||
frame_array_index = current_frame - 1
|
||||
|
||||
# Load image if available
|
||||
if frame_paths and 0 <= frame_array_index < len(frame_paths):
|
||||
frame_info = frame_paths[frame_array_index]
|
||||
image = load_image_from_disk(
|
||||
frame_info.get("filename", ""), frame_info.get("subfolder", "")
|
||||
)
|
||||
|
||||
if image is not None:
|
||||
# Resize mask to match image dimensions
|
||||
img_h, img_w = image.shape[:2]
|
||||
mask_resized = resize_mask_to_image(mask_batch[i], img_h, img_w)
|
||||
|
||||
# Apply mask overlay
|
||||
masked_image = apply_mask_overlay(image, mask_resized, background_color)
|
||||
masked_images.append(masked_image)
|
||||
raw_images.append(image)
|
||||
else:
|
||||
# Image load failed, create empty tensors matching mask dimensions
|
||||
empty_img = np.zeros((height, width, 3), dtype=np.float32)
|
||||
masked_images.append(empty_img)
|
||||
raw_images.append(empty_img)
|
||||
else:
|
||||
# No frame_paths or out of bounds, create empty tensors matching mask dimensions
|
||||
empty_img = np.zeros((height, width, 3), dtype=np.float32)
|
||||
masked_images.append(empty_img)
|
||||
raw_images.append(empty_img)
|
||||
|
||||
# Stack images into batches
|
||||
if masked_images:
|
||||
masked_batch = np.stack(masked_images, axis=0)
|
||||
masked_tensor = torch.from_numpy(masked_batch).float()
|
||||
else:
|
||||
masked_tensor = torch.zeros((frame_count, height, width, 3), dtype=torch.float32)
|
||||
|
||||
if raw_images:
|
||||
raw_batch = np.stack(raw_images, axis=0)
|
||||
raw_tensor = torch.from_numpy(raw_batch).float()
|
||||
else:
|
||||
raw_tensor = torch.zeros((frame_count, height, width, 3), dtype=torch.float32)
|
||||
|
||||
return (mask_tensor, mask_single_channel, masked_tensor, raw_tensor, curve_json)
|
||||
Reference in New Issue
Block a user