557 lines
20 KiB
Python
557 lines
20 KiB
Python
"""
|
|
Utility functions for tensor and PIL image conversions.
|
|
"""
|
|
|
|
import numpy as np
|
|
import torch
|
|
import json
|
|
import comfy.utils
|
|
from PIL import Image
|
|
from typing import List, Union, Optional
|
|
|
|
|
|
def tensor_to_pil(images: torch.Tensor) -> List[Image.Image]:
|
|
"""
|
|
Convert tensor images to PIL Images.
|
|
|
|
Args:
|
|
images: Tensor of shape [B, H, W, C] with values in [0, 1]
|
|
|
|
Returns:
|
|
List of PIL Images
|
|
|
|
Raises:
|
|
ValueError: If input is not a torch.Tensor or has unsupported shape
|
|
"""
|
|
if not isinstance(images, torch.Tensor):
|
|
raise ValueError(f"Expected torch.Tensor, got {type(images)}")
|
|
|
|
# Ensure tensor is on CPU and in correct format
|
|
images = images.cpu()
|
|
|
|
# Handle different tensor shapes
|
|
if images.dim() == 3:
|
|
# Single image [H, W, C]
|
|
images = images.unsqueeze(0)
|
|
elif images.dim() == 2:
|
|
# Grayscale [H, W]
|
|
images = images.unsqueeze(0).unsqueeze(-1)
|
|
|
|
# Convert to [0, 255] range
|
|
if images.max() <= 1.0:
|
|
images = images * 255.0
|
|
|
|
images = images.clamp(0, 255).byte()
|
|
|
|
# Convert each image in batch to PIL
|
|
pil_images = []
|
|
for img in images:
|
|
img_np = img.numpy()
|
|
if img_np.shape[-1] == 1:
|
|
# Grayscale
|
|
pil_img = Image.fromarray(img_np.squeeze(-1), mode='L')
|
|
elif img_np.shape[-1] == 3:
|
|
# RGB
|
|
pil_img = Image.fromarray(img_np, mode='RGB')
|
|
elif img_np.shape[-1] == 4:
|
|
# RGBA
|
|
pil_img = Image.fromarray(img_np, mode='RGBA')
|
|
else:
|
|
raise ValueError(f"Unsupported channel count: {img_np.shape[-1]}")
|
|
pil_images.append(pil_img)
|
|
|
|
return pil_images
|
|
|
|
|
|
def pil_to_tensor(pil_images: Union[List[Image.Image], Image.Image]) -> torch.Tensor:
|
|
"""
|
|
Convert PIL Images to tensor format.
|
|
|
|
Args:
|
|
pil_images: Single PIL Image or list of PIL Images
|
|
|
|
Returns:
|
|
Tensor of shape [B, H, W, C] with values in [0, 1]
|
|
|
|
Raises:
|
|
ValueError: If input is not a PIL Image or list of PIL Images
|
|
"""
|
|
if isinstance(pil_images, Image.Image):
|
|
pil_images = [pil_images]
|
|
|
|
if not isinstance(pil_images, list):
|
|
raise ValueError(f"Expected PIL Image or list of PIL Images, got {type(pil_images)}")
|
|
|
|
tensor_list = []
|
|
for pil_img in pil_images:
|
|
if not isinstance(pil_img, Image.Image):
|
|
raise ValueError(f"Expected PIL Image, got {type(pil_img)}")
|
|
|
|
# Convert to RGB if needed
|
|
if pil_img.mode != 'RGB':
|
|
pil_img = pil_img.convert('RGB')
|
|
|
|
# Convert to numpy array
|
|
img_np = np.array(pil_img).astype(np.float32) / 255.0
|
|
|
|
# Convert to tensor
|
|
img_tensor = torch.from_numpy(img_np)
|
|
tensor_list.append(img_tensor)
|
|
|
|
# Stack into batch
|
|
images_tensor = torch.stack(tensor_list)
|
|
|
|
return images_tensor
|
|
|
|
|
|
def masks_to_tensor(masks: Union[torch.Tensor, Image.Image, List, np.ndarray]) -> Optional[torch.Tensor]:
|
|
"""
|
|
Convert various mask formats to tensor format.
|
|
|
|
Args:
|
|
masks: Masks in various formats (torch.Tensor, Image.Image, List, np.ndarray)
|
|
|
|
Returns:
|
|
torch.Tensor [N, H, W] with values in [0, 1], or None if conversion fails
|
|
"""
|
|
if isinstance(masks, torch.Tensor):
|
|
# Ensure float type and range [0, 1]
|
|
masks = masks.float()
|
|
# Check if tensor is not empty before calling max()
|
|
if masks.numel() > 0 and masks.max() > 1.0:
|
|
masks = masks / 255.0
|
|
|
|
# Squeeze extra channel dimension if present (N, 1, H, W) -> (N, H, W)
|
|
if masks.ndim == 4 and masks.shape[1] == 1:
|
|
masks = masks.squeeze(1)
|
|
|
|
return masks.cpu()
|
|
elif isinstance(masks, np.ndarray):
|
|
masks = torch.from_numpy(masks).float()
|
|
# Check if tensor is not empty before calling max()
|
|
if masks.numel() > 0 and masks.max() > 1.0:
|
|
masks = masks / 255.0
|
|
|
|
# Squeeze extra channel dimension if present
|
|
if masks.ndim == 4 and masks.shape[1] == 1:
|
|
masks = masks.squeeze(1)
|
|
|
|
return masks
|
|
|
|
return masks
|
|
|
|
def draw_visualize_image(image, masks, scores=None, bboxs=None, alpha=0.5, stroke_width=5, font_size=24):
|
|
|
|
if isinstance(image, torch.Tensor):
|
|
# tensor_to_pil returns a list, get the first image
|
|
image = tensor_to_pil(image)[0]
|
|
elif isinstance(image, np.ndarray):
|
|
image = Image.fromarray((image * 255).astype(np.uint8) if image.max() <= 1.0 else image.astype(np.uint8))
|
|
|
|
# Convert to numpy for processing
|
|
img_np = np.array(image).astype(np.float32) / 255.0
|
|
|
|
# Resize masks to image size if needed
|
|
if isinstance(masks, torch.Tensor):
|
|
masks_np = masks.cpu().numpy()
|
|
else:
|
|
masks_np = masks
|
|
|
|
from PIL import ImageDraw, ImageFont
|
|
from scipy import ndimage
|
|
try:
|
|
font = ImageFont.load_default().font_variant(size=font_size)
|
|
except:
|
|
font = ImageFont.load_default()
|
|
|
|
# Create colored overlay
|
|
np.random.seed(42)
|
|
overlay = img_np.copy()
|
|
|
|
# Store text drawing info for later
|
|
text_info_list = []
|
|
|
|
num_masks = len(masks_np)
|
|
pbar = comfy.utils.ProgressBar(num_masks)
|
|
processed_masks = 0
|
|
for i, mask in enumerate(masks_np):
|
|
# Squeeze extra dimensions (masks may be [1, H, W] or [H, W])
|
|
while mask.ndim > 2:
|
|
mask = mask.squeeze(0)
|
|
|
|
# Resize mask to image size if needed
|
|
if mask.shape != img_np.shape[:2]:
|
|
from PIL import Image as PILImage
|
|
mask_pil = PILImage.fromarray((mask * 255).astype(np.uint8))
|
|
mask_pil = mask_pil.resize((img_np.shape[1], img_np.shape[0]), PILImage.NEAREST)
|
|
mask = np.array(mask_pil).astype(np.float32) / 255.0
|
|
|
|
# Random color for this mask
|
|
color = np.random.rand(3)
|
|
|
|
# Create darker stroke color (0.4x of original color for darker effect)
|
|
stroke_color = color * 0.4
|
|
|
|
# Create thick stroke using morphological operations
|
|
binary_mask = (mask > 0.5).astype(np.uint8)
|
|
# Dilate to get outer boundary (thick stroke with configurable width)
|
|
dilated = ndimage.binary_dilation(binary_mask, iterations=stroke_width).astype(np.float32)
|
|
# Stroke is the difference between dilated and original
|
|
stroke_mask = dilated - binary_mask
|
|
|
|
# Apply stroke (darker color with full opacity, no alpha blending)
|
|
for c in range(3):
|
|
overlay[:, :, c] = np.where(
|
|
stroke_mask > 0.5,
|
|
stroke_color[c],
|
|
overlay[:, :, c]
|
|
)
|
|
|
|
# Apply colored mask (lighter fill)
|
|
for c in range(3):
|
|
overlay[:, :, c] = np.where(
|
|
mask > 0.5,
|
|
overlay[:, :, c] * (1 - alpha) + color[c] * alpha,
|
|
overlay[:, :, c]
|
|
)
|
|
|
|
# Find the center x and top y of the mask for text placement
|
|
mask_coords = np.argwhere(mask > 0.5)
|
|
if len(mask_coords) > 0:
|
|
# Calculate x center and y top
|
|
y_top = int(mask_coords[:, 0].min())
|
|
x_center = int(mask_coords[:, 1].mean())
|
|
|
|
# Convert stroke_color to int tuple for background box
|
|
stroke_color_int = tuple((stroke_color * 255).astype(int).tolist())
|
|
|
|
# Prepare text with score if available
|
|
if scores is not None:
|
|
try:
|
|
# Handle different score formats
|
|
if isinstance(scores, torch.Tensor):
|
|
# Flatten tensor and get the i-th score
|
|
scores_flat = scores.flatten()
|
|
if i < len(scores_flat):
|
|
score = scores_flat[i].item()
|
|
text = f"id:{i} score:{score:.2f}"
|
|
else:
|
|
text = f"id:{i}"
|
|
elif isinstance(scores, (list, np.ndarray)):
|
|
score = scores[i] if isinstance(scores[i], (int, float)) else scores[i].item()
|
|
text = f"id:{i} score:{score:.2f}"
|
|
elif isinstance(scores, float):
|
|
score = scores
|
|
text = f"id:{i} score:{score:.2f}"
|
|
else:
|
|
text = f"id:{i}"
|
|
except Exception as e:
|
|
text = f"id:{i}"
|
|
print(f"Error getting score {i}: {e}")
|
|
else:
|
|
text = f"id:{i}"
|
|
|
|
# Store text info for drawing later
|
|
text_info_list.append({
|
|
'text': text,
|
|
'position': (x_center, max(0, y_top - font_size)),
|
|
'bg_color': stroke_color_int
|
|
})
|
|
|
|
# Update progress bar
|
|
processed_masks += 1
|
|
pbar.update_absolute(processed_masks, num_masks)
|
|
|
|
# Convert overlay to PIL image once after all masks are processed
|
|
result = Image.fromarray((overlay * 255).astype(np.uint8))
|
|
draw = ImageDraw.Draw(result)
|
|
|
|
# Draw all text labels
|
|
for text_info in text_info_list:
|
|
text = text_info['text']
|
|
position = text_info['position']
|
|
bg_color = text_info['bg_color']
|
|
|
|
# Get text bounding box
|
|
bbox = draw.textbbox(position, text, font=font)
|
|
# Draw background box with stroke_color
|
|
padding = 8
|
|
draw.rectangle(
|
|
[(bbox[0] - padding, bbox[1] - padding),
|
|
(bbox[2] + padding, bbox[3] + padding)],
|
|
fill=bg_color
|
|
)
|
|
# Draw text in white
|
|
draw.text(position, text, fill=(255, 255, 255), font=font)
|
|
|
|
return result
|
|
|
|
|
|
def resize_mask(mask, shape):
|
|
return torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[0], shape[1]), mode="bilinear").squeeze(1)
|
|
|
|
def join_image_with_alpha(image: torch.Tensor, alpha: torch.Tensor, invert=False):
|
|
batch_size = min(len(image), len(alpha))
|
|
out_images = []
|
|
|
|
if invert:
|
|
alpha = 1.0 - resize_mask(alpha, image.shape[1:])
|
|
else:
|
|
alpha = resize_mask(alpha, image.shape[1:])
|
|
for i in range(batch_size):
|
|
out_images.append(torch.cat((image[i][:,:,:3], alpha[i].unsqueeze(2)), dim=2))
|
|
|
|
return torch.stack(out_images),
|
|
|
|
def parse_points(points_str, image_shape=None):
|
|
"""Parse point coordinates from JSON string and validate bounds.
|
|
|
|
Supports two formats:
|
|
1. {"points": [[x, y], ...], "labels": [1, 0, ...]} - Direct format with normalized coordinates
|
|
2. [{"x": x, "y": y}, ...] - Legacy format with pixel coordinates
|
|
|
|
Converts pixel coordinates to normalized coordinates (0-1 range) if image_shape is provided.
|
|
|
|
Returns:
|
|
tuple: (points_array, count) for format 1, or (points_array, count, validation_errors) for format 2
|
|
"""
|
|
if not points_str or not points_str.strip():
|
|
return None, None, []
|
|
|
|
try:
|
|
parsed_data = json.loads(points_str)
|
|
|
|
# Check if it's the new format with "points" and "labels" keys
|
|
if isinstance(parsed_data, dict) and "points" in parsed_data:
|
|
points = parsed_data["points"]
|
|
if not points:
|
|
return None, None
|
|
return points, len(points), []
|
|
|
|
# Legacy format: list of point dictionaries
|
|
if not isinstance(parsed_data, list):
|
|
raise ValueError(f"Points must be a JSON array or object with 'points' key, got {type(parsed_data).__name__}")
|
|
|
|
if len(parsed_data) == 0:
|
|
return None, None, []
|
|
|
|
points = []
|
|
validation_errors = []
|
|
|
|
for i, point_dict in enumerate(parsed_data):
|
|
if not isinstance(point_dict, dict):
|
|
err = f"Point {i} is not a dictionary"
|
|
print(f"Warning: {err}, skipping")
|
|
validation_errors.append(err)
|
|
continue
|
|
|
|
if 'x' not in point_dict or 'y' not in point_dict:
|
|
err = f"Point {i} missing 'x' or 'y' key"
|
|
print(f"Warning: {err}, skipping")
|
|
validation_errors.append(err)
|
|
continue
|
|
|
|
try:
|
|
x = float(point_dict['x'])
|
|
y = float(point_dict['y'])
|
|
|
|
# Validate coordinates are non-negative
|
|
if x < 0 or y < 0:
|
|
err = f"Point {i} has negative coordinates ({x}, {y})"
|
|
print(f"Warning: {err}, skipping")
|
|
validation_errors.append(err)
|
|
continue
|
|
|
|
# Normalize to 0-1 range if image shape is provided
|
|
if image_shape is not None:
|
|
height, width = image_shape[1], image_shape[2] # [batch, height, width, channels]
|
|
|
|
# Validate within image bounds
|
|
if x >= width or y >= height:
|
|
err = f"Point {i} ({x}, {y}) outside image bounds ({width}x{height})"
|
|
print(f"Warning: {err}, skipping")
|
|
validation_errors.append(err)
|
|
continue
|
|
|
|
# Normalize coordinates to [0, 1] range
|
|
x = x / width
|
|
y = y / height
|
|
|
|
points.append([x, y])
|
|
|
|
except (ValueError, TypeError) as e:
|
|
err = f"Could not convert point {i} coordinates to float: {e}"
|
|
print(f"Warning: {err}, skipping")
|
|
validation_errors.append(err)
|
|
continue
|
|
|
|
if not points:
|
|
return None, None, validation_errors
|
|
|
|
return points, len(points), validation_errors
|
|
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f"Invalid JSON in points: {str(e)}")
|
|
except Exception as e:
|
|
print(f"Error parsing points: {e}")
|
|
return None, None, [str(e)]
|
|
|
|
def parse_bbox(bbox, image_shape=None):
|
|
"""Parse bounding box from BBOX type (tuple/list/dict) and validate
|
|
|
|
Supports multiple formats:
|
|
1. {"boxes": [[x, y, w, h], ...], "labels": [true/false, ...]} - Direct format with normalized coordinates
|
|
2. KJNodes: [{'startX': x, 'startY': y, 'endX': x2, 'endY': y2}, ...]
|
|
3. Tuple/list: (x1, y1, x2, y2) or (x, y, width, height)
|
|
4. Dict: {'startX': x, 'startY': y, 'endX': x2, 'endY': y2}
|
|
|
|
Converts pixel coordinates to normalized coordinates (0-1 range) if image_shape is provided.
|
|
|
|
Returns:
|
|
tuple: (boxes_array, count) for all formats
|
|
"""
|
|
if bbox is None:
|
|
return None, 0
|
|
|
|
try:
|
|
# Check if it's a string that needs to be parsed as JSON
|
|
if isinstance(bbox, str):
|
|
bbox = json.loads(bbox)
|
|
|
|
# Check if it's the new format with "boxes" and "labels" keys
|
|
if isinstance(bbox, dict) and "boxes" in bbox:
|
|
boxes = bbox["boxes"]
|
|
if not boxes:
|
|
return None, 0
|
|
return boxes, len(boxes)
|
|
|
|
all_coords = []
|
|
|
|
# Try to extract coordinates regardless of type checks
|
|
# This handles cases where ComfyUI wraps data in unexpected ways
|
|
if hasattr(bbox, '__iter__') and not isinstance(bbox, (str, bytes)):
|
|
# It's some kind of sequence
|
|
try:
|
|
bbox_list = list(bbox)
|
|
if len(bbox_list) == 0:
|
|
return None, 0
|
|
|
|
# Check if it's a list of 4 numbers (single bbox)
|
|
if len(bbox_list) == 4 and all(isinstance(x, (int, float)) for x in bbox_list):
|
|
coords = [float(x) for x in bbox_list]
|
|
all_coords.append(coords)
|
|
else:
|
|
# Process each element as a potential bbox
|
|
for elem in bbox_list:
|
|
coords = None
|
|
|
|
# Try to access as dict-like (KJNodes format)
|
|
if hasattr(elem, '__getitem__'):
|
|
try:
|
|
x1 = float(elem['startX'])
|
|
y1 = float(elem['startY'])
|
|
x2 = float(elem['endX'])
|
|
y2 = float(elem['endY'])
|
|
coords = [x1, y1, x2, y2]
|
|
except (KeyError, TypeError):
|
|
# Not dict format, might be numeric sequence
|
|
pass
|
|
|
|
# If still no coords, try as numeric sequence
|
|
if coords is None:
|
|
if hasattr(elem, '__iter__') and not isinstance(elem, (str, bytes)):
|
|
inner = list(elem)
|
|
if len(inner) == 4:
|
|
coords = [float(x) for x in inner]
|
|
if coords is not None:
|
|
all_coords.append(coords)
|
|
|
|
except Exception as e:
|
|
raise ValueError(f"Failed to process bbox as sequence: {e}")
|
|
|
|
# Try single dict format
|
|
elif hasattr(bbox, '__getitem__'):
|
|
try:
|
|
x1 = float(bbox['startX'])
|
|
y1 = float(bbox['startY'])
|
|
x2 = float(bbox['endX'])
|
|
y2 = float(bbox['endY'])
|
|
coords = [x1, y1, x2, y2]
|
|
all_coords.append(coords)
|
|
except (KeyError, TypeError) as e:
|
|
raise ValueError(f"Dictionary bbox missing required keys: {e}")
|
|
|
|
else:
|
|
raise ValueError(f"Unsupported bbox type: {type(bbox)}")
|
|
|
|
if not all_coords:
|
|
raise ValueError(
|
|
f"Could not extract coordinates from bbox. Type: {type(bbox)}, Content: {repr(bbox)[:200]}")
|
|
|
|
# Process and validate each bbox
|
|
validated_coords = []
|
|
for coords in all_coords:
|
|
# Handle xywh format (convert to xyxy)
|
|
x1, y1, x2, y2 = coords
|
|
if x2 < x1 or y2 < y1:
|
|
# Assume xywh format: (x, y, width, height)
|
|
width, height = x2, y2
|
|
x2 = x1 + width
|
|
y2 = y1 + height
|
|
coords = [x1, y1, x2, y2]
|
|
|
|
# Validate coordinates
|
|
if coords[0] >= coords[2]:
|
|
raise ValueError(f"Invalid bbox: x1 ({coords[0]}) must be < x2 ({coords[2]})")
|
|
if coords[1] >= coords[3]:
|
|
raise ValueError(f"Invalid bbox: y1 ({coords[1]}) must be < y2 ({coords[3]})")
|
|
if coords[0] < 0 or coords[1] < 0:
|
|
raise ValueError(f"Bounding box coordinates must be non-negative, got x1={coords[0]}, y1={coords[1]}")
|
|
|
|
# Normalize to 0-1 range if image shape is provided
|
|
if image_shape is not None:
|
|
height, width = image_shape[1], image_shape[2]
|
|
|
|
# Validate within image bounds
|
|
if coords[0] >= width or coords[2] > width:
|
|
print(f"Warning: bbox x coordinates ({coords[0]}, {coords[2]}) outside image width ({width})")
|
|
if coords[1] >= height or coords[3] > height:
|
|
print(f"Warning: bbox y coordinates ({coords[1]}, {coords[3]}) outside image height ({height})")
|
|
|
|
# Normalize coordinates to [0, 1] range
|
|
x1 = coords[0] / width
|
|
y1 = coords[1] / height
|
|
x2 = coords[2] / width
|
|
y2 = coords[3] / height
|
|
new_coords = [
|
|
(x1 + x2) / 2,
|
|
(y1 + y2) / 2,
|
|
x2 - x1,
|
|
y2 - y1
|
|
]
|
|
|
|
validated_coords.append(new_coords)
|
|
else:
|
|
validated_coords.append(coords)
|
|
|
|
return validated_coords, len(validated_coords)
|
|
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f"Invalid JSON in bbox: {str(e)}")
|
|
except (ValueError, TypeError) as e:
|
|
error_msg = f"Invalid bbox: {str(e)}\n"
|
|
error_msg += f"Input type: {type(bbox)}\n"
|
|
error_msg += f"Input content: {repr(bbox)[:500]}"
|
|
raise ValueError(error_msg)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
bboxes = [
|
|
[
|
|
159.9,
|
|
189.5,
|
|
317.4,
|
|
329.3
|
|
]
|
|
]
|
|
print(parse_bbox(bboxes, image_shape=(1, 832, 480, 3))) |