Spaces:
Sleeping
Sleeping
File size: 3,399 Bytes
2cf467c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """Composite utilities for combining node results."""
import numpy as np
from typing import List, Tuple
from PIL import Image
def composite_nodes(nodes: List[dict], bboxes: List[Tuple[float, float, float, float]],
container_bbox: Tuple[float, float, float, float]) -> Tuple[np.ndarray, np.ndarray]:
"""Composite multiple nodes into a single mask.
Args:
nodes: List of node dictionaries with "mask" key
bboxes: List of (x, y, w, h) bounding boxes for each node
container_bbox: Container bounding box (x, y, w, h)
Returns:
Tuple of (composite_mask, composite_sdf)
- composite_mask: Combined mask (H, W)
- composite_sdf: Combined SDF (H, W) - simplified version
"""
# Handle container_bbox format
if isinstance(container_bbox, (tuple, list)) and len(container_bbox) == 4:
cx, cy, cw, ch = container_bbox
elif isinstance(container_bbox, dict):
cx = container_bbox.get("x", 0)
cy = container_bbox.get("y", 0)
cw = container_bbox.get("width", container_bbox.get("w", 0))
ch = container_bbox.get("height", container_bbox.get("h", 0))
else:
raise ValueError(f"Invalid container_bbox format: {container_bbox}")
cw_int = int(float(cw))
ch_int = int(float(ch))
# Initialize composite mask
composite_mask = np.zeros((ch_int, cw_int), dtype=np.float32)
# Place each node's mask at its bbox position
for i, (node, bbox) in enumerate(zip(nodes, bboxes)):
mask = node.get("mask")
if mask is None:
continue
# Handle different bbox formats
if isinstance(bbox, (tuple, list)) and len(bbox) == 4:
x, y, w, h = bbox
elif isinstance(bbox, dict):
x = bbox.get("x", 0)
y = bbox.get("y", 0)
w = bbox.get("width", bbox.get("w", 0))
h = bbox.get("height", bbox.get("h", 0))
else:
print(f"Warning: Invalid bbox format at index {i}: {bbox} (type: {type(bbox)})")
continue
x_int = int(float(x))
y_int = int(float(y))
w_int = int(float(w))
h_int = int(float(h))
# Resize mask to bbox size if needed
if mask.shape != (h_int, w_int):
from scipy.ndimage import zoom
zoom_y = h_int / mask.shape[0]
zoom_x = w_int / mask.shape[1]
mask_resized = zoom(mask, (zoom_y, zoom_x), order=1)
else:
mask_resized = mask
# Clip to valid range
mask_resized = np.clip(mask_resized, 0, 1)
# Place mask in composite
y_end = min(y_int + h_int, ch_int)
x_end = min(x_int + w_int, cw_int)
y_start = max(0, y_int)
x_start = max(0, x_int)
if y_end > y_start and x_end > x_start:
mask_crop = mask_resized[:y_end-y_start, :x_end-x_start]
composite_mask[y_start:y_end, x_start:x_end] = np.maximum(
composite_mask[y_start:y_end, x_start:x_end],
mask_crop
)
# Generate simplified SDF from composite mask
# This is a simplified version - full SDF generation would use distance_transform_edt
composite_sdf = composite_mask - 0.5 # Simple approximation
return composite_mask, composite_sdf
|