Spaces:
Sleeping
Sleeping
File size: 12,304 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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 | """Save hierarchical layout optimization results."""
import numpy as np
from PIL import Image, ImageDraw
from typing import Dict, Any, Tuple, Optional
import os
def save_hierarchical_result(result: Dict[str, Any], save_path: str = "hierarchical_result.png",
base_dir: str = ".") -> None:
"""Save hierarchical optimization result as an image.
Args:
result: Result dictionary from HierarchicalOptimizer.optimize_tree()
save_path: Path to save the final image
base_dir: Base directory for resolving image paths
"""
# Get root bbox
root_bbox = result.get("final_bbox")
if root_bbox is None:
print("Warning: No final_bbox found in result")
return
# Handle different bbox formats
if isinstance(root_bbox, (tuple, list)) and len(root_bbox) == 4:
x, y, w, h = root_bbox
elif isinstance(root_bbox, dict):
x = root_bbox.get("x", 0)
y = root_bbox.get("y", 0)
w = root_bbox.get("width", root_bbox.get("w", 1000))
h = root_bbox.get("height", root_bbox.get("h", 1000))
else:
print(f"Warning: Invalid root_bbox format: {root_bbox}")
return
root_x = int(float(x))
root_y = int(float(y))
Wc = int(float(w))
Hc = int(float(h))
# First pass: calculate the actual bounding box needed for all nodes
# Root node's final_bbox is absolute, so pass 0,0 as offset (will use final_bbox directly)
min_x, min_y, max_x, max_y = _calculate_bounds(result, offset_x=0, offset_y=0, is_root=True)
# Add some padding to ensure nothing is cut off
padding = 10
canvas_width = max_x - min_x + padding * 2
canvas_height = max_y - min_y + padding * 2
canvas_offset_x = min_x - padding
canvas_offset_y = min_y - padding
# Create canvas
canvas = Image.new("RGBA", (canvas_width, canvas_height), (255, 255, 255, 255))
# Print root bbox info
print(f"\n[Saving hierarchical result]")
print(f" Root bbox: ({root_x}, {root_y}, {Wc}, {Hc})")
print(f" Content bounds: ({min_x}, {min_y}) to ({max_x}, {max_y})")
print(f" Canvas size: {canvas_width}x{canvas_height} (offset: {canvas_offset_x}, {canvas_offset_y})")
# Second pass: recursively composite all nodes
# Root node's final_bbox is absolute, so we need to adjust for canvas offset
# For root node, we pass a flag indicating it's the root
_composite_node_to_canvas(result, canvas, base_dir,
offset_x=root_x - canvas_offset_x,
offset_y=root_y - canvas_offset_y,
is_root=True)
# Convert to RGB and save
canvas_rgb = Image.new("RGB", canvas.size, (255, 255, 255))
canvas_rgb.paste(canvas, mask=canvas.split()[3])
canvas_rgb.save(save_path, "PNG")
print(f"Hierarchical layout result saved to: {save_path}\n")
def _draw_bbox(canvas: Image.Image, x: int, y: int, w: int, h: int, node_type: str) -> None:
"""Draw bounding box on canvas.
Args:
canvas: PIL Image canvas to draw on
x: X coordinate (absolute)
y: Y coordinate (absolute)
w: Width
h: Height
node_type: Type of node (for color selection)
"""
# Clip coordinates to canvas bounds
x_clip = max(0, min(x, canvas.width - 1))
y_clip = max(0, min(y, canvas.height - 1))
x_end = min(x + w, canvas.width)
y_end = min(y + h, canvas.height)
if x_end <= x_clip or y_end <= y_clip:
return
# Choose color based on node type
color_map = {
"column": (255, 0, 0, 255), # Red for column
"row": (0, 255, 0, 255), # Green for row
"layer": (0, 0, 255, 255), # Blue for layer
"chart": (255, 165, 0, 255), # Orange for chart
"image": (255, 0, 255, 255), # Magenta for image
"text": (0, 255, 255, 255), # Cyan for text
}
color = color_map.get(node_type, (128, 128, 128, 255)) # Gray for unknown types
# Draw rectangle
draw = ImageDraw.Draw(canvas)
draw.rectangle([x_clip, y_clip, x_end - 1, y_end - 1], outline=color, width=2)
def _calculate_bounds(node_result: Dict[str, Any], offset_x: int = 0, offset_y: int = 0,
is_root: bool = False) -> Tuple[int, int, int, int]:
"""Calculate the bounding box of all nodes in the tree.
Args:
node_result: Node result dictionary
offset_x: X offset accumulated from parent containers (for relative coordinates)
offset_y: Y offset accumulated from parent containers (for relative coordinates)
is_root: Whether this is the root node (root's final_bbox is absolute, others are relative)
Returns:
Tuple of (min_x, min_y, max_x, max_y) in absolute coordinates
"""
# Get node bbox
bbox = node_result.get("final_bbox")
if bbox is None:
return (0, 0, 0, 0)
# 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:
return (0, 0, 0, 0)
# Root node's bbox is absolute, other nodes' bboxes are relative to parent container
if is_root:
# Root bbox is already absolute, use it directly
x_abs = int(float(x))
y_abs = int(float(y))
else:
# Convert relative coordinates to absolute by adding parent offset
x_abs = offset_x + int(float(x))
y_abs = offset_y + int(float(y))
w_int = max(1, int(float(w)))
h_int = max(1, int(float(h)))
# Initialize bounds with this node's bounds
min_x = x_abs
min_y = y_abs
max_x = x_abs + w_int
max_y = y_abs + h_int
# Recursively process children
# Children are not root nodes, so pass is_root=False
children = node_result.get("children", [])
for child_result in children:
child_min_x, child_min_y, child_max_x, child_max_y = _calculate_bounds(
child_result, offset_x=x_abs, offset_y=y_abs, is_root=False
)
if child_min_x < min_x:
min_x = child_min_x
if child_min_y < min_y:
min_y = child_min_y
if child_max_x > max_x:
max_x = child_max_x
if child_max_y > max_y:
max_y = child_max_y
return (min_x, min_y, max_x, max_y)
def _composite_node_to_canvas(node_result: Dict[str, Any], canvas: Image.Image,
base_dir: str, offset_x: int = 0, offset_y: int = 0,
is_root: bool = False) -> None:
"""Recursively composite node results onto canvas.
Args:
node_result: Node result dictionary
canvas: PIL Image canvas to composite onto
base_dir: Base directory for resolving image paths
offset_x: X offset accumulated from parent containers (for canvas offset adjustment)
offset_y: Y offset accumulated from parent containers (for canvas offset adjustment)
is_root: Whether this is the root node (root's final_bbox is absolute, others are relative)
"""
# Get node bbox
bbox = node_result.get("final_bbox")
if bbox is None:
return
# 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:
return
# Root node's bbox is absolute, other nodes' bboxes are relative to parent container
if is_root:
# Root bbox is already absolute, offset_x/offset_y here are canvas offsets (root_x - canvas_offset_x)
# For root: final_bbox x is root_x (absolute), offset_x = root_x - canvas_offset_x
# We want canvas-relative: root_x - canvas_offset_x = offset_x
# So we use offset_x directly (it's already the canvas-relative position)
x_abs = offset_x
y_abs = offset_y
else:
# Child node bbox is relative to parent container
# Convert to absolute coordinates by adding parent offset
x_abs = offset_x + int(float(x))
y_abs = offset_y + int(float(y))
w_int = max(1, int(float(w)))
h_int = max(1, int(float(h)))
# Get node type for debugging
node_type = node_result.get("type", "unknown")
# Check if this is a leaf node with an image
metadata = node_result.get("metadata", {})
image_path = metadata.get("image_path") or metadata.get("full_path") or node_result.get("image_path")
# Try to resolve path
if image_path:
if os.path.isabs(image_path):
full_path = image_path
else:
full_path = os.path.join(base_dir, image_path)
if os.path.exists(full_path):
# Load and place image
img = Image.open(full_path).convert("RGBA")
# Resize image while preserving aspect ratio
# Calculate scale to fit within bbox
img_w, img_h = img.size
scale_w = w_int / img_w if img_w > 0 else 1.0
scale_h = h_int / img_h if img_h > 0 else 1.0
scale = min(scale_w, scale_h) # Use smaller scale to fit within bbox
# Calculate new size preserving aspect ratio
new_w = int(img_w * scale)
new_h = int(img_h * scale)
# Resize image
img_resized = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
# Center image within bbox
x_offset = (w_int - new_w) // 2
y_offset = (h_int - new_h) // 2
# Calculate final position on canvas
x_final = x_abs + x_offset
y_final = y_abs + y_offset
# Clip coordinates to canvas bounds
x_clip = max(0, min(x_final, canvas.width - 1))
y_clip = max(0, min(y_final, canvas.height - 1))
# Calculate how much of the image fits
x_end = min(x_clip + new_w, canvas.width)
y_end = min(y_clip + new_h, canvas.height)
w_fit = x_end - x_clip
h_fit = y_end - y_clip
if w_fit > 0 and h_fit > 0:
if w_fit < new_w or h_fit < new_h:
img_resized = img_resized.crop((0, 0, w_fit, h_fit))
canvas.paste(img_resized, (x_clip, y_clip), img_resized)
# Print bbox info for debugging
coord_type = "absolute" if is_root else "relative"
print(f" [Save] Node '{node_type}': image={os.path.basename(image_path)}, "
f"bbox=({int(float(x))}, {int(float(y))}, {w_int}, {h_int}) [{coord_type}], "
f"img_size=({img_w}x{img_h}→{new_w}x{new_h}), "
f"placed_at=({x_final}, {y_final}) [canvas-relative]")
else:
# Print bbox info even if image doesn't exist
coord_type = "absolute" if is_root else "relative"
print(f" [Save] Node '{node_type}': bbox=({int(float(x))}, {int(float(y))}, {w_int}, {h_int}) [{coord_type}], "
f"image_path={image_path} (not found)")
else:
# Print bbox info for nodes without image_path (container nodes, text nodes, etc.)
coord_type = "absolute" if is_root else "relative"
print(f" [Save] Node '{node_type}': bbox=({int(float(x))}, {int(float(y))}, {w_int}, {h_int}) [{coord_type}], no image")
# Draw bounding box
_draw_bbox(canvas, x_abs, y_abs, w_int, h_int, node_type)
# Recursively process children
# Children's coordinates are relative to this container, so pass this container's absolute position as offset
# Children are not root nodes
children = node_result.get("children", [])
for child_result in children:
_composite_node_to_canvas(child_result, canvas, base_dir, offset_x=x_abs, offset_y=y_abs, is_root=False)
|