"""Hierarchical layout optimizer with extensible architecture.""" from typing import Dict, List, Optional, Tuple, Any import numpy as np from dataclasses import dataclass, field from . import parameters as params from .constraints import ( ConstraintProcessor, RelativeSizeProcessor, PaddingProcessor, OrientationProcessor, OverlapProcessor, AlignmentProcessor, GapProcessor, ) from .handlers import ( NodeHandler, ImageNodeHandler, TextNodeHandler, ) from .strategies import OptimizationStrategy, SDFOptimizationStrategy, RuleBasedLayoutStrategy from .utils.parser import parse_layout_tree, LayoutNode from .utils.composite import composite_nodes from .utils.placeholder import create_placeholder_rectangle, create_placeholder_rounded_rectangle @dataclass class OptimizationConfig: """Optimization configuration.""" strategy: OptimizationStrategy = field(default_factory=SDFOptimizationStrategy) constraint_processors: Dict[str, ConstraintProcessor] = field(default_factory=dict) node_handlers: Dict[str, NodeHandler] = field(default_factory=dict) weights: Dict[str, float] = field(default_factory=dict) optimization_params: Dict[str, Any] = field(default_factory=dict) placeholder_config: Dict[str, Any] = field(default_factory=dict) base_dir: Optional[str] = None device: Optional[str] = None debug: bool = False # Enable debug mode: visualization and detailed logging use_rule_based: bool = True # Enable rule-based layout for row/column nodes (faster) rule_based_types: List[str] = field(default_factory=lambda: ['row', 'column']) # Types that use rule-based layout def __post_init__(self): """Initialize default values.""" if not self.constraint_processors: self.constraint_processors = { "relative_size": RelativeSizeProcessor(), "padding": PaddingProcessor(), "orientation": OrientationProcessor(), "overlap": OverlapProcessor(), "gap": GapProcessor(), } if not self.node_handlers: self.node_handlers = { "image": ImageNodeHandler(), "chart": ImageNodeHandler(), "text": TextNodeHandler(), } if not self.weights: self.weights = { "w_similarity": params.W_SIMILARITY, "w_readability": params.W_READABILITY, "w_alignment_consistency": params.W_ALIGNMENT_CONSISTENCY, "w_alignment_similarity": params.W_ALIGNMENT_SIMILARITY, "w_proximity": params.W_PROXIMITY, } if not self.optimization_params: self.optimization_params = { "opt_res_list": params.OPT_RES_LIST, "outer_rounds": params.OUTER_ROUNDS, "inner_steps": params.INNER_STEPS, "lr": params.LEARNING_RATE, } class HierarchicalOptimizer: """Hierarchical layout optimizer with extensible architecture.""" def __init__(self, config: Optional[OptimizationConfig] = None): """Initialize optimizer. Args: config: Optimization configuration (uses default if None) """ self.config = config or OptimizationConfig() self._register_default_processors() self._register_default_handlers() def _register_default_processors(self): """Register default constraint processors.""" if not self.config.constraint_processors: self.config.constraint_processors = { "relative_size": RelativeSizeProcessor(), "padding": PaddingProcessor(), "orientation": OrientationProcessor(), "overlap": OverlapProcessor(), } def _register_default_handlers(self): """Register default node handlers.""" if not self.config.node_handlers: self.config.node_handlers = { "image": ImageNodeHandler(), "chart": ImageNodeHandler(), "text": TextNodeHandler(), } def register_processor(self, processor: ConstraintProcessor): """Register a custom constraint processor. Args: processor: Constraint processor instance """ # Register for all constraint types it can handle for constraint_type in ["relative_size", "padding", "orientation", "overlap", "gap", "alignment"]: if processor.can_handle(constraint_type): self.config.constraint_processors[constraint_type] = processor def register_handler(self, handler: NodeHandler): """Register a custom node handler. Args: handler: Node handler instance """ # Register for all node types it can handle for node_type in ["image", "chart", "text", "shape", "layer", "column", "row"]: if handler.can_handle(node_type): self.config.node_handlers[node_type] = handler def set_strategy(self, strategy: OptimizationStrategy): """Set optimization strategy. Args: strategy: Optimization strategy instance """ self.config.strategy = strategy def optimize_tree(self, tree_json: dict) -> Dict[str, Any]: """Optimize entire tree structure. Args: tree_json: JSON dictionary with layout tree structure Returns: Dictionary with optimization results for each node """ # Parse tree root_node = parse_layout_tree(tree_json) # print("Root node: ", root_node) # Get container bbox from root root_bbox = ( root_node.bbox.get("x", 0), root_node.bbox.get("y", 0), root_node.bbox.get("width", 1000), root_node.bbox.get("height", 1000), ) # Optimize recursively from bottom up result = self._optimize_node(root_node, root_bbox, "root") return result def _optimize_node(self, node: LayoutNode, parent_bbox: Tuple[float, float, float, float], node_path: str = "root") -> Dict[str, Any]: """Optimize a single node and its children recursively. Args: node: Layout node to optimize parent_bbox: Parent container bounding box (x, y, w, h) node_path: Path string for this node (e.g., "root.child0.child1") Returns: Dictionary with optimization results """ result = { "type": node.type, "bbox": node.bbox, "final_bbox": None, "image_path": getattr(node, "image_path", None), # Preserve image_path for saving } # Check if node is a container (has children) if node.children: # Check if we can use rule-based layout (much faster for row/column) if self._can_use_rule_based_layout(node): return self._rule_based_layout(node, parent_bbox, node_path) # Container node: optimize children using SDF optimization # print(f"[HierarchicalOptimizer] Processing container node: type={node.type}, num_children={len(node.children)}, path={node_path}") child_results = [] child_nodes_data = [] # First, recursively optimize all children for i, child in enumerate(node.children): child_path = f"{node_path}.child{i}" child_result = self._optimize_node(child, parent_bbox, child_path) child_results.append(child_result) # Load child node data (masks, etc.) # Use final_bbox from child_result if available (for container nodes that have been optimized), # otherwise use initial bbox for i, child in enumerate(node.children): child_result = child_results[i] # Get bbox: use final_bbox from child_result if child is a container that has been optimized if child_result.get("final_bbox") is not None: # Convert final_bbox tuple to dict format for consistency final_bbox = child_result["final_bbox"] if isinstance(final_bbox, (tuple, list)) and len(final_bbox) >= 4: child_bbox = { "x": final_bbox[0], "y": final_bbox[1], "width": final_bbox[2], "height": final_bbox[3] } else: child_bbox = child.bbox else: child_bbox = child.bbox # Check if child has composite_mask (from previous optimization) # Use composite_mask if available, as it represents the actual shape mask = None metadata = {} if child_result.get("composite_mask") is not None: mask = child_result["composite_mask"] # print(f"[HierarchicalOptimizer] Using composite_mask for child {i} (type={child.type})") else: handler = self._get_handler(child.type) if handler: mask, metadata = handler.load(child.__dict__, self.config.base_dir) else: # Fallback: use placeholder bbox = child_bbox width = bbox.get("width", 100) if isinstance(bbox, dict) else (bbox[2] if isinstance(bbox, (tuple, list)) else 100) height = bbox.get("height", 100) if isinstance(bbox, dict) else (bbox[3] if isinstance(bbox, (tuple, list)) else 100) # For container nodes (layer/column/row), create a rounded rectangle instead of solid rectangle if child.type in ["layer", "column", "row"]: try: _, mask = create_placeholder_rounded_rectangle(width, height) except Exception: # Fallback to rectangle if rounded rectangle is unavailable _, mask = create_placeholder_rectangle(width, height) else: _, mask = create_placeholder_rectangle(width, height) metadata = {"placeholder": True} child_nodes_data.append({ "mask": mask, "bbox": child_bbox, "metadata": metadata, "type": child.type, }) # Get container bbox (use node's bbox or parent's) container_bbox = ( node.bbox.get("x", 0), node.bbox.get("y", 0), node.bbox.get("width", parent_bbox[2]), node.bbox.get("height", parent_bbox[3]), ) # Optimize children layout constraints = node.constraints or {} # Generate unique save prefix for this node using node path # Replace dots and special characters to make valid filename save_prefix = node_path.replace(".", "_").replace(" ", "_") # Extract grandchildren info for proximity ratio calculation # For each child, collect its children's bboxes (grandchildren of current container) grandchildren_list = [] for child_result in child_results: grandchildren = [] # Check if child_result has children (from recursive optimization) child_children = child_result.get("children", []) for grandchild_result in child_children: grandchild_bbox = grandchild_result.get("final_bbox") if grandchild_bbox: if isinstance(grandchild_bbox, (tuple, list)) and len(grandchild_bbox) >= 4: grandchildren.append(tuple(grandchild_bbox[:4])) elif isinstance(grandchild_bbox, dict): grandchildren.append(( grandchild_bbox.get("x", 0), grandchild_bbox.get("y", 0), grandchild_bbox.get("width", grandchild_bbox.get("w", 0)), grandchild_bbox.get("height", grandchild_bbox.get("h", 0)) )) grandchildren_list.append(grandchildren) config = { **self.config.optimization_params, "device": self.config.device, "debug": self.config.debug, # Pass debug flag "w_similarity": self.config.weights.get("w_similarity", 1.0), "w_readability": self.config.weights.get("w_readability", 1.0), "w_alignment_consistency": self.config.weights.get("w_alignment_consistency", 1.0), "w_alignment_similarity": self.config.weights.get("w_alignment_similarity", params.W_ALIGNMENT_SIMILARITY), "w_proximity": self.config.weights.get("w_proximity", params.W_PROXIMITY), "container_type": node.type, # Pass container type (column, row, or layer) "grandchildren_list": grandchildren_list, # Pass grandchildren for proximity calculation } # print(f"[HierarchicalOptimizer] Calling strategy.optimize for container_type={node.type}, num_children={len(child_nodes_data)}") # print("child_nodes_data:", child_nodes_data) # print("config:", config) optimized_bboxes = self.config.strategy.optimize( child_nodes_data, container_bbox, constraints, config, save_prefix=save_prefix, ) # print(f"[HierarchicalOptimizer] Strategy returned optimized_bboxes: {optimized_bboxes}") # Calculate actual container bbox based on children's layout results # Note: optimized_bboxes are relative to container origin (0,0) if optimized_bboxes: # Convert bbox to tuple format if needed (handle both dict and tuple formats) def bbox_to_tuple(bbox): if isinstance(bbox, dict): return ( bbox.get("x", 0), bbox.get("y", 0), bbox.get("width", bbox.get("w", 0)), bbox.get("height", bbox.get("h", 0)) ) elif isinstance(bbox, (tuple, list)) and len(bbox) >= 4: return tuple(bbox[:4]) else: return (0, 0, 0, 0) bbox_tuples = [bbox_to_tuple(bbox) for bbox in optimized_bboxes] # print("bbox_tuples:", bbox_tuples) # Find the bounding box that contains all children (for position adjustment) min_x = min(bbox[0] for bbox in bbox_tuples) min_y = min(bbox[1] for bbox in bbox_tuples) # Adjust children's bboxes: shift them so that min_x and min_y become 0 # This makes children relative to the new container origin (0,0) # For container nodes, preserve their own calculated width/height adjusted_bboxes_for_composite = [] for i, (child_result, opt_bbox) in enumerate(zip(child_results, optimized_bboxes)): bbox_tuple = bbox_to_tuple(opt_bbox) # Adjust coordinates: subtract min_x and min_y to start from (0,0) adjusted_x = bbox_tuple[0] - min_x adjusted_y = bbox_tuple[1] - min_y # For container nodes, preserve the width/height calculated from their children # Only update position (x, y), not size (w, h) if child_result.get("children") and child_result.get("final_bbox"): # Child is a container with its own final_bbox calculated from its children child_final = child_result["final_bbox"] if isinstance(child_final, (tuple, list)) and len(child_final) >= 4: # Use the child's own calculated width and height adjusted_bbox = ( adjusted_x, adjusted_y, child_final[2], # Keep child's calculated width child_final[3] # Keep child's calculated height ) else: adjusted_bbox = (adjusted_x, adjusted_y, bbox_tuple[2], bbox_tuple[3]) else: # Leaf node or no children: use optimized bbox size adjusted_bbox = (adjusted_x, adjusted_y, bbox_tuple[2], bbox_tuple[3]) adjusted_bboxes_for_composite.append(adjusted_bbox) # Store relative coordinates (relative to parent container) in final_bbox # save_result.py will convert to absolute coordinates by adding parent offsets child_result["final_bbox"] = adjusted_bbox if i < len(node.children): node.children[i].final_bbox = adjusted_bbox # print(f"node.children[{i}].final_bbox:", node.children[i].final_bbox) # Calculate actual container size based on adjusted bboxes # This ensures the container size is based on children's preserved widths/heights actual_max_x = max(bbox[0] + bbox[2] for bbox in adjusted_bboxes_for_composite) actual_max_y = max(bbox[1] + bbox[3] for bbox in adjusted_bboxes_for_composite) actual_width = actual_max_x actual_height = actual_max_y # print(f"actual_container_size: width={actual_width:.2f}, height={actual_height:.2f}") # Container position: keep original container position # Container size: actual size needed to contain all children # Children are now adjusted to start from (0,0) relative to container actual_container_bbox = ( container_bbox[0], # Keep original x position container_bbox[1], # Keep original y position actual_width, # Actual width needed actual_height # Actual height needed ) else: # Fallback to original container bbox if no children actual_container_bbox = container_bbox adjusted_bboxes_for_composite = optimized_bboxes # Update child results with optimized bboxes (no adjustment needed) for i, (child_result, opt_bbox) in enumerate(zip(child_results, optimized_bboxes)): child_result["final_bbox"] = opt_bbox if i < len(node.children): node.children[i].final_bbox = opt_bbox # Composite children results (use actual container bbox and adjusted bboxes) composite_mask, composite_sdf = self.config.strategy.composite( child_nodes_data, adjusted_bboxes_for_composite, actual_container_bbox, ) result["final_bbox"] = actual_container_bbox result["composite_mask"] = composite_mask result["composite_sdf"] = composite_sdf result["children"] = child_results else: # Leaf node: just load data handler = self._get_handler(node.type) if handler: mask, metadata = handler.load(node.__dict__, self.config.base_dir) result["mask"] = mask result["metadata"] = metadata result["final_bbox"] = ( node.bbox.get("x", 0), node.bbox.get("y", 0), node.bbox.get("width", 100), node.bbox.get("height", 100), ) return result def _get_handler(self, node_type: str) -> Optional[NodeHandler]: """Get handler for node type. Args: node_type: Type of node Returns: Node handler or None """ return self.config.node_handlers.get(node_type) def _can_use_rule_based_layout(self, node: LayoutNode) -> bool: """Check if node can use rule-based layout instead of SDF optimization. Rule-based layout is much faster and more accurate for simple row/column layouts. Args: node: Layout node to check Returns: True if rule-based layout can be used """ # Check if rule-based is enabled if not self.config.use_rule_based: return False # Check if node type is in the allowed list if node.type not in self.config.rule_based_types: return False # Need at least 2 children to benefit from rule-based layout if not node.children or len(node.children) < 2: return False # Check for complex overlap constraints that require SDF constraints = node.constraints or {} if 'overlap' in constraints: # Overlap constraints need precise collision detection - use SDF return False return True def _rule_based_layout(self, node: LayoutNode, parent_bbox: Tuple[float, float, float, float], node_path: str) -> Dict[str, Any]: """Execute rule-based layout for simple row/column arrangements. Args: node: Layout node to optimize parent_bbox: Parent container bounding box (x, y, w, h) node_path: Path string for this node Returns: Dictionary with optimization results """ # print(f"[HierarchicalOptimizer] Using rule-based layout for {node.type} node: {node_path}") result = { "type": node.type, "bbox": node.bbox, "final_bbox": None, "image_path": getattr(node, "image_path", None), } # First, recursively optimize all children child_results = [] for i, child in enumerate(node.children): child_path = f"{node_path}.child{i}" child_result = self._optimize_node(child, parent_bbox, child_path) child_results.append(child_result) # Load child node data child_nodes_data = [] for i, child in enumerate(node.children): child_result = child_results[i] # Get bbox from child result if child_result.get("final_bbox") is not None: final_bbox = child_result["final_bbox"] if isinstance(final_bbox, (tuple, list)) and len(final_bbox) >= 4: child_bbox = { "x": final_bbox[0], "y": final_bbox[1], "width": final_bbox[2], "height": final_bbox[3] } else: child_bbox = child.bbox else: child_bbox = child.bbox # Get mask mask = None metadata = {} if child_result.get("composite_mask") is not None: mask = child_result["composite_mask"] else: handler = self._get_handler(child.type) if handler: mask, metadata = handler.load(child.__dict__, self.config.base_dir) else: # Fallback placeholder bbox = child_bbox width = bbox.get("width", 100) if isinstance(bbox, dict) else bbox[2] height = bbox.get("height", 100) if isinstance(bbox, dict) else bbox[3] if child.type in ["layer", "column", "row"]: _, mask = create_placeholder_rounded_rectangle(width, height) else: _, mask = create_placeholder_rectangle(width, height) metadata = {"placeholder": True} child_nodes_data.append({ "mask": mask, "bbox": child_bbox, "metadata": metadata, "type": child.type, }) # Get container bbox container_bbox = ( node.bbox.get("x", 0), node.bbox.get("y", 0), node.bbox.get("width", parent_bbox[2]), node.bbox.get("height", parent_bbox[3]), ) # Use rule-based strategy from .strategies import RuleBasedLayoutStrategy rule_strategy = RuleBasedLayoutStrategy() constraints = node.constraints or {} config = { "container_type": node.type, } # Calculate optimized bboxes using rule-based layout optimized_bboxes = rule_strategy.optimize( child_nodes_data, container_bbox, constraints, config, ) # Adjust bboxes to start from (0, 0) relative to container if optimized_bboxes: min_x = min(bbox[0] for bbox in optimized_bboxes) min_y = min(bbox[1] for bbox in optimized_bboxes) adjusted_bboxes_for_composite = [] for i, (child_result, opt_bbox) in enumerate(zip(child_results, optimized_bboxes)): adjusted_x = opt_bbox[0] - min_x adjusted_y = opt_bbox[1] - min_y # For container nodes, preserve their calculated size if child_result.get("children") and child_result.get("final_bbox"): child_final = child_result["final_bbox"] if isinstance(child_final, (tuple, list)) and len(child_final) >= 4: adjusted_bbox = ( adjusted_x, adjusted_y, child_final[2], # Keep child's calculated width child_final[3] # Keep child's calculated height ) else: adjusted_bbox = (adjusted_x, adjusted_y, opt_bbox[2], opt_bbox[3]) else: adjusted_bbox = (adjusted_x, adjusted_y, opt_bbox[2], opt_bbox[3]) adjusted_bboxes_for_composite.append(adjusted_bbox) child_result["final_bbox"] = adjusted_bbox if i < len(node.children): node.children[i].final_bbox = adjusted_bbox # Calculate actual container size actual_max_x = max(bbox[0] + bbox[2] for bbox in adjusted_bboxes_for_composite) actual_max_y = max(bbox[1] + bbox[3] for bbox in adjusted_bboxes_for_composite) actual_width = actual_max_x actual_height = actual_max_y actual_container_bbox = ( container_bbox[0], container_bbox[1], actual_width, actual_height ) else: actual_container_bbox = container_bbox adjusted_bboxes_for_composite = optimized_bboxes for i, (child_result, opt_bbox) in enumerate(zip(child_results, optimized_bboxes)): child_result["final_bbox"] = opt_bbox if i < len(node.children): node.children[i].final_bbox = opt_bbox # Composite children results composite_mask, composite_sdf = rule_strategy.composite( child_nodes_data, adjusted_bboxes_for_composite, actual_container_bbox, ) result["final_bbox"] = actual_container_bbox result["composite_mask"] = composite_mask result["composite_sdf"] = composite_sdf result["children"] = child_results return result