"""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)