"""Main optimization function for SDF-based layout optimization.""" import os import shutil from typing import List, Tuple, Optional, Dict import numpy as np import torch import torch.nn.functional as F from modules.infographics_generator.layout_system import parameters as params from .core import ( load_binary_mask_from_rgba, tight_bbox_ratio, binary_to_sdf_norm, dilate_mask, make_container_grid, sdf_to_softmask, area_sum, ) from .bbox import ( bbox_aspect_from_unconstrained, unconstrained_from_bbox, ) from .losses import ( compute_alignment_consistency_loss, compute_alignment_similarity_loss, compute_readability_loss, compute_proximity_ratio_loss, compute_visual_balance_loss, compute_position_size_similarity_loss, ) from .visualization import ( visualize_optimization_progress, visualize_final_result, save_composite_image, ) def optimize( png_list, # List of image paths for N nodes original_png_list=None, # List of original image paths for saving composite Wc=1000, Hc=1000, opt_res_list=params.OPT_RES_LIST, # optimize on these resolutions outer_rounds=params.OUTER_ROUNDS, # augmented-lagrangian outer updates per stage inner_steps=params.INNER_STEPS, # gradient steps per outer round tau_schedule=params.TAU_SCHEDULE, rho_init=params.RHO_INIT, # initial penalty parameter rho_mult=params.RHO_MULT, # penalty multiplier lr=params.LEARNING_RATE, size_min=params.SIZE_MIN, # Legacy parameter for backward compatibility min_sizes=None, # List of (min_width, min_height) tuples for each element pen_weight=params.PEN_WEIGHT, # weight for penetration penalty pen_eta_px=params.PEN_ETA_PX, reference_bboxes=None, # List of reference element bboxes (x, y, w, h) from Example layout reference_parent_bbox=None, # Reference parent container bbox (x, y, w, h) w_similarity=params.W_SIMILARITY, # weight for position/size similarity loss size_rules=None, # List of tuples (source_idx, target_idx) for size hierarchy rules w_readability=params.W_READABILITY, # weight for readability loss w_alignment_consistency=params.W_ALIGNMENT_CONSISTENCY, # weight for alignment consistency loss alignment_constraint=None, # Dictionary with alignment constraint from JSON (direction, value) w_alignment_similarity=params.W_ALIGNMENT_SIMILARITY, # weight for alignment similarity loss proximity_info=None, # Dictionary with container hierarchy info for proximity ratio calculation w_proximity=params.W_PROXIMITY, # weight for proximity ratio loss w_data_ink=params.W_DATA_INK, # weight for data ink loss (maximize union area) w_visual_balance=params.W_VISUAL_BALANCE, # weight for visual balance loss min_gap_px=20.0, # minimum gap between elements in pixels device=None, save_prefix=None, # Prefix for saving result images (None = use default names) debug=False, # Enable debug mode: visualization and detailed logging ): """Main SDF-based layout optimization function. Args: png_list: List of image paths for N nodes original_png_list: List of original image paths for saving composite Wc, Hc: Container width and height opt_res_list: Optimize on these resolutions outer_rounds: Augmented-lagrangian outer updates per stage inner_steps: Gradient steps per outer round tau_schedule: Schedule for tau parameter rho_init: Initial penalty parameter rho_mult: Penalty multiplier lr: Learning rate size_min: Legacy minimum size parameter min_sizes: List of (min_width, min_height) tuples for each element pen_weight: Weight for penetration penalty pen_eta_px: Eta parameter for penetration penalty reference_bboxes: List of reference element bboxes from Example layout reference_parent_bbox: Reference parent container bbox w_similarity: Weight for position/size similarity loss size_rules: List of tuples (source_idx, target_idx) for size hierarchy rules w_readability: Weight for readability loss w_alignment_consistency: Weight for alignment consistency loss alignment_constraint: Dictionary with alignment constraint from JSON w_alignment_similarity: Weight for alignment similarity loss proximity_info: Dictionary with container hierarchy info w_proximity: Weight for proximity ratio loss w_data_ink: Weight for data ink loss w_visual_balance: Weight for visual balance loss min_gap_px: Minimum gap between elements in pixels device: PyTorch device save_prefix: Prefix for saving result images debug: Enable debug mode (visualization and detailed logging) Returns: List of final bounding boxes [(x, y, w, h), ...] """ if device is None: device = "cuda" if torch.cuda.is_available() else "cpu" print("Device:", device) # Validate png_list if png_list is None or len(png_list) < 1: raise ValueError("png_list must be provided with at least one image path") num_nodes = len(png_list) # Handle original_png_list if original_png_list is None: original_png_list = png_list # Set dilation_radii to 100 for all elements dilation_radii = [20.0] * num_nodes # print(f"Number of nodes: {num_nodes}") # print(f"Dilation radii: {dilation_radii}") # Validate reference bboxes if reference_bboxes is None: reference_bboxes = [] if reference_parent_bbox is None: reference_parent_bbox = (0.0, 0.0, float(Wc), float(Hc)) # Normalize reference bboxes: subtract x_min and y_min to make coordinates start from 0 if reference_bboxes and len(reference_bboxes) > 0: # Find minimum x and y across all reference bboxes x_min = min(bbox[0] for bbox in reference_bboxes) y_min = min(bbox[1] for bbox in reference_bboxes) # Subtract x_min and y_min from all bboxes normalized_reference_bboxes = [] for bbox in reference_bboxes: if isinstance(bbox, (tuple, list)) and len(bbox) >= 4: normalized_bbox = (bbox[0] - x_min, bbox[1] - y_min, bbox[2], bbox[3]) elif isinstance(bbox, dict): normalized_bbox = { 'x': bbox.get('x', 0) - x_min, 'y': bbox.get('y', 0) - y_min, 'width': bbox.get('width', bbox.get('w', 100)), 'height': bbox.get('height', bbox.get('h', 100)) } else: normalized_bbox = bbox normalized_reference_bboxes.append(normalized_bbox) reference_bboxes = normalized_reference_bboxes # Also normalize reference_parent_bbox if reference_parent_bbox: x_p, y_p, w_p, h_p = reference_parent_bbox reference_parent_bbox = (x_p - x_min, y_p - y_min, w_p, h_p) # print(f"Normalized reference bboxes: subtracted x_min={x_min:.1f}, y_min={y_min:.1f}") # Validate size rules if size_rules is None: size_rules = [] # Validate and set up min_sizes if min_sizes is None: # Use default values for all elements min_sizes = [(params.MIN_WIDTH_DEFAULT, params.MIN_HEIGHT_DEFAULT)] * num_nodes elif len(min_sizes) < num_nodes: # Extend with default values default_size = (params.MIN_WIDTH_DEFAULT, params.MIN_HEIGHT_DEFAULT) min_sizes = min_sizes + [default_size] * (num_nodes - len(min_sizes)) elif len(min_sizes) > num_nodes: # Truncate to num_nodes min_sizes = min_sizes[:num_nodes] # print(f"Reference bboxes: {len(reference_bboxes)} elements") # print(f"Reference parent bbox: {reference_parent_bbox}") # print(f"Size rules: {len(size_rules)} rules") # print(f"Similarity weight: {w_similarity}, Readability weight: {w_readability}") # print(f"Min sizes: {min_sizes}") # print(f"Debug mode: {debug}") # Create visualization folder for this optimization run (only in debug mode) viz_folder = None if debug: if save_prefix: viz_folder = f"{save_prefix}_visualization" else: viz_folder = "optimization_visualization" # Clear existing folder contents if it exists if os.path.exists(viz_folder): shutil.rmtree(viz_folder) os.makedirs(viz_folder, exist_ok=True) # print(f"Visualization folder created: {viz_folder}") # load masks for all nodes masks = [] ratios = [] for i, png_path in enumerate(png_list): if png_path is None: raise ValueError(f"Image path at index {i} is None") mask = load_binary_mask_from_rgba(png_path) # Apply dilation with radius 100 to all elements mask = dilate_mask(mask, dilation_radii[i]) # print(f"Applied dilation to node {i} mask with radius {dilation_radii[i]:.1f} pixels") masks.append(mask) r, _ = tight_bbox_ratio(mask) ratios.append(r) # print(f"Node {i} aspect ratio (tight alpha bbox): r={r:.4f}") # precompute SDF templates for all nodes sdf_norms = [] sdf_tensors = [] for i, mask in enumerate(masks): sdf_norm = binary_to_sdf_norm(mask, pad=16) sdf_norms.append(sdf_norm) sdf_t = torch.from_numpy(sdf_norm)[None, None].to(device) sdf_tensors.append(sdf_t) # learnable parameters: (tx,ty,ts) for each object # Initialize from reference_bboxes if available, otherwise use zeros opt_params = [] for i in range(num_nodes): if reference_bboxes and i < len(reference_bboxes): # Initialize from reference bbox ref_bbox = reference_bboxes[i] if isinstance(ref_bbox, (tuple, list)) and len(ref_bbox) >= 4: ref_x, ref_y, ref_w, ref_h = ref_bbox[0], ref_bbox[1], ref_bbox[2], ref_bbox[3] elif isinstance(ref_bbox, dict): ref_x = ref_bbox.get("x", 0) ref_y = ref_bbox.get("y", 0) ref_w = ref_bbox.get("width", ref_bbox.get("w", 100)) ref_h = ref_bbox.get("height", ref_bbox.get("h", 100)) else: ref_x, ref_y, ref_w, ref_h = 0, 0, 100, 100 # Convert reference bbox to unconstrained parameters # Use material's aspect ratio (ratios[i]) instead of JSON's w/h min_width, min_height = min_sizes[i] tx_init, ty_init, ts_init = unconstrained_from_bbox( ref_x, ref_y, ref_w, ref_h, Wc, Hc, ratios[i], min_width=min_width, min_height=min_height, size_min=size_min ) # print(f"Node {i}: Initializing from reference bbox ({ref_x:.1f}, {ref_y:.1f}, {ref_w:.1f}, {ref_h:.1f}), " # f"adjusted to aspect ratio {ratios[i]:.4f}") else: # Initialize with zeros (default) tx_init, ty_init, ts_init = 0.0, 0.0, 0.0 tx = torch.nn.Parameter(torch.tensor(tx_init, device=device)) ty = torch.nn.Parameter(torch.tensor(ty_init, device=device)) ts = torch.nn.Parameter(torch.tensor(ts_init, device=device)) opt_params.extend([tx, ty, ts]) opt = torch.optim.Adam(opt_params, lr=lr) # augmented lagrangian multipliers for constraint g = A_inter = 0 lam = torch.tensor(0.0, device=device) rho = torch.tensor(rho_init, device=device) tau_list = list(tau_schedule) if len(tau_list) < outer_rounds: tau_list += [tau_list[-1]] * (outer_rounds - len(tau_list)) # staged optimization over resolutions for stage_idx, stage_res in enumerate(opt_res_list): Hs = Ws = int(stage_res) X, Y = make_container_grid(Hs, Ws, device=device) # print(f"\n=== Stage optimize at {Ws}x{Hs} (container {Wc}x{Hc}) ===") # Visualize initial state before optimization (only for first stage and only in debug mode) if stage_idx == 0 and debug: with torch.no_grad(): # Use full resolution for visualization X_init, Y_init = make_container_grid(Hc, Wc, device=device) initial_bboxes = [] initial_softmasks = [] for i in range(num_nodes): tx_idx = i * 3 ty_idx = i * 3 + 1 ts_idx = i * 3 + 2 tx = opt_params[tx_idx] ty = opt_params[ty_idx] ts = opt_params[ts_idx] min_width, min_height = min_sizes[i] x, y, w, h = bbox_aspect_from_unconstrained( tx, ty, ts, Wc, Hc, ratios[i], min_width=min_width, min_height=min_height, size_min=size_min ) # print(f"Node {i} initial bbox: ({x.item():.1f}, {y.item():.1f}, {w.item():.1f}, {h.item():.1f})") initial_bboxes.append((x.item(), y.item(), w.item(), h.item())) # Use smaller tau for initial visualization to show actual mask shape m, _ = sdf_to_softmask(sdf_tensors[i], x, y, w, h, X_init, Y_init, tau_px=0.5) initial_softmasks.append(m) # Compute initial loss values for display initial_union = torch.ones_like(initial_softmasks[0]) for m in initial_softmasks: initial_union = initial_union * (1.0 - m) initial_union = 1.0 - initial_union initial_A_union = initial_union.sum().item() # Full resolution, da=1 initial_inter = torch.zeros_like(initial_softmasks[0]) for i in range(num_nodes): for j in range(i + 1, num_nodes): initial_inter = initial_inter + initial_softmasks[i] * initial_softmasks[j] initial_A_inter = initial_inter.sum().item() # Full resolution, da=1 # Compute visual balance loss for initial state (using bboxes) initial_bboxes_tensors = [] for i in range(num_nodes): tx_idx = i * 3 ty_idx = i * 3 + 1 ts_idx = i * 3 + 2 tx = opt_params[tx_idx] ty = opt_params[ty_idx] ts = opt_params[ts_idx] min_width, min_height = min_sizes[i] x, y, w, h = bbox_aspect_from_unconstrained( tx, ty, ts, Wc, Hc, ratios[i], min_width=min_width, min_height=min_height, size_min=size_min ) initial_bboxes_tensors.append((x, y, w, h)) L_visual_balance_init = compute_visual_balance_loss( initial_bboxes_tensors, Wc, Hc, device=device ) initial_loss_info = { 'A_union': initial_A_union, 'A_inter': initial_A_inter, 'visual_balance': w_visual_balance * L_visual_balance_init.item(), } initial_save_path = os.path.join(viz_folder, f"initial_stage{stage_idx}.png") visualize_optimization_progress( initial_softmasks, initial_bboxes, Wc, Hc, epoch=-1, loss_info=initial_loss_info, save_path=initial_save_path ) for k in range(outer_rounds): tau_px = float(tau_list[min(k, len(tau_list)-1)]) for t in range(inner_steps): opt.zero_grad(set_to_none=True) # Compute bboxes for all nodes bboxes = [] softmasks = [] distances = [] for i in range(num_nodes): tx_idx = i * 3 ty_idx = i * 3 + 1 ts_idx = i * 3 + 2 tx = opt_params[tx_idx] ty = opt_params[ty_idx] ts = opt_params[ts_idx] min_width, min_height = min_sizes[i] x, y, w, h = bbox_aspect_from_unconstrained( tx, ty, ts, Wc, Hc, ratios[i], min_width=min_width, min_height=min_height, size_min=size_min ) bboxes.append((x, y, w, h)) m, d_px = sdf_to_softmask(sdf_tensors[i], x, y, w, h, X, Y, tau_px=tau_px) softmasks.append(m) distances.append(d_px) # Compute union: 1 - product of (1 - mask_i) union = torch.ones_like(softmasks[0]) for m in softmasks: union = union * (1.0 - m) union = 1.0 - union # Compute intersection: sum of all pairwise intersections inter = torch.zeros_like(softmasks[0]) for i in range(num_nodes): for j in range(i + 1, num_nodes): inter = inter + softmasks[i] * softmasks[j] A_union = area_sum(union, Wc, Hc) A_inter = area_sum(inter, Wc, Hc) # must go to 0 # Visual balance loss (using bboxes directly for differentiability) L_visual_balance = compute_visual_balance_loss( bboxes, Wc, Hc, device=device ) # Penetration loss: based on actual bbox gap (not limited by SDF range) # Penalize when gap < min_gap_px (ensures minimum gap between elements) # Use ReLU for hard cutoff: no penalty when gap >= min_gap_px # For "layer" container type, skip this loss (overlapping is allowed) L_pen = torch.tensor(0.0, device=device) # # Check if this is a layer container (overlapping allowed) # is_layer = False # if proximity_info: # container_types = proximity_info.get("types", []) # if len(container_types) > 0: # is_layer = (container_types[0] == "layer") # if not is_layer: # # Only apply penetration penalty for row/column layouts # for i in range(num_nodes): # for j in range(i + 1, num_nodes): # x_i, y_i, w_i, h_i = bboxes[i] # x_j, y_j, w_j, h_j = bboxes[j] # # Compute gap in each dimension (negative if overlapping) # gap_x = torch.max(x_j - (x_i + w_i), x_i - (x_j + w_j)) # gap_y = torch.max(y_j - (y_i + h_i), y_i - (y_j + h_j)) # # Combined gap logic: # # - If both separated (both positive): Euclidean distance # # - If both overlapping (both negative): use MAX (smallest overlap, easiest to fix) # # - If one separated, one overlapping: use the separated one (that's the actual gap) # if gap_x >= 0 and gap_y >= 0: # # Both separated: Euclidean distance # gap = torch.sqrt(gap_x * gap_x + gap_y * gap_y) # elif gap_x < 0 and gap_y < 0: # # Both overlapping: use the smaller overlap (easier to separate) # gap = torch.max(gap_x, gap_y) # else: # # One separated, one overlapping: they're aligned in overlapping dimension # # Use the separated dimension's gap # gap = torch.max(gap_x, gap_y) # # Penalty if gap < min_gap_px # L_pen = L_pen + F.relu(min_gap_px - gap) # Position/Size similarity loss L_similarity = torch.tensor(0.0, device=device) if reference_bboxes and len(reference_bboxes) >= num_nodes: # Current generated bboxes as tensors generated_bboxes = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes] # Generated parent container bbox (current container) generated_parent_bbox = (0.0, 0.0, float(Wc), float(Hc)) L_similarity = compute_position_size_similarity_loss( reference_bboxes[:num_nodes], generated_bboxes, reference_parent_bbox, generated_parent_bbox, device=device ) # Readability loss (size hierarchy consistency) L_readability = torch.tensor(0.0, device=device) if size_rules and len(size_rules) > 0: generated_bboxes_readability = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes] L_readability = compute_readability_loss( size_rules, generated_bboxes_readability, size_ratio_threshold=params.SIZE_RATIO_THRESHOLD, device=device ) # Alignment consistency loss (hierarchical alignment) L_alignment_consistency = torch.tensor(0.0, device=device) if reference_bboxes and len(reference_bboxes) >= num_nodes: generated_bboxes_alignment = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes] generated_parent_bbox_alignment = (0.0, 0.0, float(Wc), float(Hc)) L_alignment_consistency = compute_alignment_consistency_loss( reference_bboxes[:num_nodes], generated_bboxes_alignment, reference_parent_bbox, generated_parent_bbox_alignment, device=device ) # Alignment similarity loss (based on JSON constraint) L_alignment_similarity = torch.tensor(0.0, device=device) if alignment_constraint and w_alignment_similarity > 0: generated_bboxes_alignment_sim = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes] container_bbox_alignment = (0.0, 0.0, float(Wc), float(Hc)) L_alignment_similarity = compute_alignment_similarity_loss( generated_bboxes_alignment_sim, container_bbox_alignment, alignment_constraint, device=device ) # Proximity ratio loss L_proximity = torch.tensor(0.0, device=device) if proximity_info and w_proximity > 0: # Extract information from proximity_info container_bboxes = proximity_info.get("containers", []) child_bboxes_list = proximity_info.get("children", []) grandchild_bboxes_list = proximity_info.get("grandchildren", []) container_types = proximity_info.get("types", []) container_weights = proximity_info.get("weights", None) # Convert current generated bboxes to tuples for proximity calculation # For N-element case: treat as single container with N children if len(container_bboxes) == 0: # Simplified N-element case: create a container with N children # Note: container type should be provided in proximity_info container_bbox = (0.0, 0.0, float(Wc), float(Hc)) # Keep bboxes as tensors for gradient computation child_bboxes = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes] # For N-element case, grandchildren would be empty (children are leaf nodes) grandchild_bboxes = [] container_type = container_types[0] if container_types else "row" # Default to row L_proximity = compute_proximity_ratio_loss( [container_bbox], [child_bboxes], [grandchild_bboxes], [container_type], container_weights=[1.0] if container_weights is None else container_weights, epsilon=params.PROXIMITY_EPSILON, device=device ) else: # Use provided proximity_info # Keep bboxes as tensors for gradient computation generated_bboxes_proximity = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes] # Update child_bboxes_list with current generated bboxes if needed # This is a simplified approach - in full implementation, we'd need to map # generated bboxes to the correct containers if len(child_bboxes_list) > 0 and len(child_bboxes_list[0]) == num_nodes: # Update first container's children with generated bboxes (as tensors) updated_child_bboxes_list = [generated_bboxes_proximity] + child_bboxes_list[1:] else: updated_child_bboxes_list = child_bboxes_list L_proximity = compute_proximity_ratio_loss( container_bboxes, updated_child_bboxes_list, grandchild_bboxes_list, container_types, container_weights, epsilon=params.PROXIMITY_EPSILON, device=device ) g = A_inter # Data ink loss: maximize union area (minimize white space) # Negative because we want to maximize A_union (minimize -A_union) L_data_ink = -A_union # Scale pen_weight by rho to prevent L_pen from being overwhelmed when rho is large # When rho is large, AL constraint dominates, so we need to scale pen_weight accordingly pen_weight_scaled = pen_weight * (1.0 + rho.item() / 1e4) # AL constraint on overlap + penalty term + similarity loss + readability loss + alignment consistency loss + alignment similarity loss + proximity loss + data ink loss + visual balance loss loss = (lam * g + 0.5 * rho * g * g + pen_weight_scaled * L_pen + w_similarity * L_similarity + w_readability * L_readability + w_alignment_consistency * L_alignment_consistency + w_alignment_similarity * L_alignment_similarity + w_proximity * L_proximity + w_data_ink * L_data_ink + w_visual_balance * L_visual_balance) loss.backward() opt.step() # outer AL update with torch.no_grad(): # Compute bboxes for logging and AL update bboxes_log = [] for i in range(num_nodes): tx_idx = i * 3 ty_idx = i * 3 + 1 ts_idx = i * 3 + 2 tx = opt_params[tx_idx] ty = opt_params[ty_idx] ts = opt_params[ts_idx] min_width, min_height = min_sizes[i] x, y, w, h = bbox_aspect_from_unconstrained( tx, ty, ts, Wc, Hc, ratios[i], min_width=min_width, min_height=min_height, size_min=size_min ) bboxes_log.append((x, y, w, h)) # Compute A_inter for AL update (always needed) softmasks_log = [] for i in range(num_nodes): x, y, w, h = bboxes_log[i] m, d_px = sdf_to_softmask(sdf_tensors[i], x, y, w, h, X, Y, tau_px=tau_px) softmasks_log.append(m) inter_log = torch.zeros_like(softmasks_log[0]) for i in range(num_nodes): for j in range(i + 1, num_nodes): inter_log = inter_log + softmasks_log[i] * softmasks_log[j] A_inter = area_sum(inter_log, Wc, Hc) # Update Lagrangian multiplier lam = lam + rho * A_inter rho = rho * rho_mult # Detailed logging and visualization (only in debug mode) if debug: # Compute all loss components for detailed logging distances_log = [] for i in range(num_nodes): x, y, w, h = bboxes_log[i] m, d_px = sdf_to_softmask(sdf_tensors[i], x, y, w, h, X, Y, tau_px=tau_px) distances_log.append(d_px) union_log = torch.ones_like(softmasks_log[0]) for m in softmasks_log: union_log = union_log * (1.0 - m) union_log = 1.0 - union_log A_union = area_sum(union_log, Wc, Hc) A_union = area_sum(union_log, Wc, Hc) # Recompute penalty for logging (based on actual bbox gap) # Skip for layer containers (overlapping allowed) L_pen_val = torch.tensor(0.0, device=device) # Check if this is a layer container is_layer = False if proximity_info: container_types = proximity_info.get("types", []) if len(container_types) > 0: is_layer = (container_types[0] == "layer") if not is_layer: for i in range(num_nodes): for j in range(i + 1, num_nodes): x_i, y_i, w_i, h_i = bboxes_log[i] x_j, y_j, w_j, h_j = bboxes_log[j] # Compute gap in each dimension gap_x = torch.max(x_j - (x_i + w_i), x_i - (x_j + w_j)) gap_y = torch.max(y_j - (y_i + h_i), y_i - (y_j + h_j)) # Combined gap logic if gap_x >= 0 and gap_y >= 0: # Both separated: Euclidean distance gap = torch.sqrt(gap_x * gap_x + gap_y * gap_y) elif gap_x < 0 and gap_y < 0: # Both overlapping: use smaller overlap gap = torch.max(gap_x, gap_y) else: # One separated, one overlapping gap = torch.max(gap_x, gap_y) # Penalty if gap < min_gap_px L_pen_val = L_pen_val + F.relu(min_gap_px - gap) # Recompute similarity loss for logging L_similarity_val = torch.tensor(0.0, device=device) if reference_bboxes and len(reference_bboxes) >= num_nodes: generated_bboxes_log = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes_log] generated_parent_bbox_log = (0.0, 0.0, float(Wc), float(Hc)) L_similarity_val = compute_position_size_similarity_loss( reference_bboxes[:num_nodes], generated_bboxes_log, reference_parent_bbox, generated_parent_bbox_log, device=device ) # Recompute readability loss for logging L_readability_val = torch.tensor(0.0, device=device) if size_rules and len(size_rules) > 0: generated_bboxes_readability_log = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes_log] L_readability_val = compute_readability_loss( size_rules, generated_bboxes_readability_log, size_ratio_threshold=params.SIZE_RATIO_THRESHOLD, device=device ) # Recompute alignment consistency loss for logging L_alignment_consistency_val = torch.tensor(0.0, device=device) if reference_bboxes and len(reference_bboxes) >= num_nodes: generated_bboxes_alignment_log = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes_log] generated_parent_bbox_alignment_log = (0.0, 0.0, float(Wc), float(Hc)) L_alignment_consistency_val = compute_alignment_consistency_loss( reference_bboxes[:num_nodes], generated_bboxes_alignment_log, reference_parent_bbox, generated_parent_bbox_alignment_log, device=device ) # Recompute alignment similarity loss for logging L_alignment_similarity_val = torch.tensor(0.0, device=device) if alignment_constraint and w_alignment_similarity > 0: generated_bboxes_alignment_sim_log = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes_log] container_bbox_alignment_log = (0.0, 0.0, float(Wc), float(Hc)) L_alignment_similarity_val = compute_alignment_similarity_loss( generated_bboxes_alignment_sim_log, container_bbox_alignment_log, alignment_constraint, device=device ) # Recompute proximity ratio loss for logging L_proximity_val = torch.tensor(0.0, device=device) if proximity_info and w_proximity > 0: container_bboxes = proximity_info.get("containers", []) child_bboxes_list = proximity_info.get("children", []) grandchild_bboxes_list = proximity_info.get("grandchildren", []) container_types = proximity_info.get("types", []) container_weights = proximity_info.get("weights", None) if len(container_bboxes) == 0: container_bbox = (0.0, 0.0, float(Wc), float(Hc)) child_bboxes = [(bbox[0].item(), bbox[1].item(), bbox[2].item(), bbox[3].item()) for bbox in bboxes_log] grandchild_bboxes = [] container_type = container_types[0] if container_types else "row" # Default to row L_proximity_val = compute_proximity_ratio_loss( [container_bbox], [child_bboxes], [grandchild_bboxes], [container_type], container_weights=[1.0] if container_weights is None else container_weights, epsilon=params.PROXIMITY_EPSILON, device=device ) else: generated_bboxes_proximity_log = [(bbox[0].item(), bbox[1].item(), bbox[2].item(), bbox[3].item()) for bbox in bboxes_log] if len(child_bboxes_list) > 0 and len(child_bboxes_list[0]) == num_nodes: updated_child_bboxes_list_log = [generated_bboxes_proximity_log] + child_bboxes_list[1:] else: updated_child_bboxes_list_log = child_bboxes_list L_proximity_val = compute_proximity_ratio_loss( container_bboxes, updated_child_bboxes_list_log, grandchild_bboxes_list, container_types, container_weights, epsilon=params.PROXIMITY_EPSILON, device=device ) # Compute data ink loss for logging L_data_ink_val = -A_union # Compute visual balance loss for logging (using bboxes_log) L_visual_balance_val = compute_visual_balance_loss( bboxes_log, Wc, Hc, device=device ) # Compute total loss for visualization g_val = A_inter total_loss_val = (lam.item() * g_val.item() + 0.5 * rho.item() * g_val.item() * g_val.item() + pen_weight * L_pen_val.item() + w_similarity * L_similarity_val.item() + w_readability * L_readability_val.item() + w_alignment_consistency * L_alignment_consistency_val.item() + w_alignment_similarity * L_alignment_similarity_val.item() + w_proximity * L_proximity_val.item() + w_data_ink * L_data_ink_val.item() + w_visual_balance * L_visual_balance_val.item()) print(f"[outer {k:02d}] tau={tau_px:.3f} A_union={A_union.item():.2f} A_inter={A_inter.item():.6f} " f"L_pen={pen_weight*L_pen_val.item():.4f} L_sim={w_similarity*L_similarity_val.item():.4f} " f"L_read={w_readability*L_readability_val.item():.4f} " f"L_align_cons={w_alignment_consistency*L_alignment_consistency_val.item():.4f} " f"L_align_sim={w_alignment_similarity*L_alignment_similarity_val.item():.4f} " f"L_prox={w_proximity*L_proximity_val.item():.4f} " f"L_data_ink={w_data_ink*L_data_ink_val.item():.4f} " f"L_balance={w_visual_balance*L_visual_balance_val.item():.4f} " f"lam={lam.item():.3e} rho={rho.item():.3e}") # Visualize optimization progress at end of each outer epoch # Use full resolution for visualization X_viz, Y_viz = make_container_grid(Hc, Wc, device=device) epoch_bboxes = [] epoch_softmasks = [] for i, bbox_log in enumerate(bboxes_log): x, y, w, h = bbox_log epoch_bboxes.append((x.item(), y.item(), w.item(), h.item())) m_viz, _ = sdf_to_softmask(sdf_tensors[i], x, y, w, h, X_viz, Y_viz, tau_px=tau_px) epoch_softmasks.append(m_viz) epoch_loss_info = { 'total': total_loss_val, 'A_union': A_union.item(), 'A_inter': A_inter.item(), 'pen': L_pen_val.item(), 'similarity': w_similarity * L_similarity_val.item(), 'readability': w_readability * L_readability_val.item(), 'alignment': w_alignment_consistency * L_alignment_consistency_val.item() + w_alignment_similarity * L_alignment_similarity_val.item(), 'alignment_consistency': w_alignment_consistency * L_alignment_consistency_val.item(), 'alignment_similarity': w_alignment_similarity * L_alignment_similarity_val.item(), 'proximity': w_proximity * L_proximity_val.item(), 'data_ink': w_data_ink * L_data_ink_val.item(), } epoch_save_path = os.path.join(viz_folder, f"stage{stage_idx}_epoch{k:02d}.png") visualize_optimization_progress( epoch_softmasks, epoch_bboxes, Wc, Hc, epoch=k, loss_info=epoch_loss_info, save_path=epoch_save_path ) else: # Non-debug mode: simple logging print(f"[outer {k:02d}] tau={tau_px:.3f} A_inter={A_inter.item():.6f} lam={lam.item():.3e} rho={rho.item():.3e}") # final bbox (continuous) final_bboxes = [] with torch.no_grad(): for i in range(num_nodes): tx_idx = i * 3 ty_idx = i * 3 + 1 ts_idx = i * 3 + 2 tx = opt_params[tx_idx] ty = opt_params[ty_idx] ts = opt_params[ts_idx] min_width, min_height = min_sizes[i] x, y, w, h = bbox_aspect_from_unconstrained( tx, ty, ts, Wc, Hc, ratios[i], min_width=min_width, min_height=min_height, size_min=size_min ) final_bboxes.append((x.item(), y.item(), w.item(), h.item())) # hard evaluation at full 1000x1000: overlap using (SDF<0) AND with torch.no_grad(): Xf, Yf = make_container_grid(Hc, Wc, device=device) # use a small tau for union display (not needed for hard overlap) softmasks_f = [] distances_f = [] for i in range(num_nodes): x, y, w, h = final_bboxes[i] m, d_px = sdf_to_softmask(sdf_tensors[i], torch.tensor(x, device=device), torch.tensor(y, device=device), torch.tensor(w, device=device), torch.tensor(h, device=device), Xf, Yf, tau_px=0.2) softmasks_f.append(m) distances_f.append(d_px) # Compute union union_f = torch.ones_like(softmasks_f[0]) for m in softmasks_f: union_f = union_f * (1.0 - m) union_f = 1.0 - union_f A_union_f = union_f.sum().item() # da=1 at full res # hard inside test: d_px < 0 for all pairs hard_overlap = torch.zeros_like(softmasks_f[0]) for i in range(num_nodes): for j in range(i + 1, num_nodes): overlap_ij = ((distances_f[i] < 0.0) & (distances_f[j] < 0.0)).float() hard_overlap = hard_overlap + overlap_ij A_overlap_hard = hard_overlap.sum().item() # print("\n=== Final Results ===") # for i, bbox in enumerate(final_bboxes): # print(f"bbox{i+1} (x,y,w,h) = {bbox}") # print(f"Union area (approx, {Wc}x{Hc}) = {A_union_f:.2f} -> ratio {A_union_f/(Wc*Hc):.4f}") # print(f"Hard overlap area (SDF<0) = {A_overlap_hard:.0f} pixels") # Save visualization and composite images (only in debug mode) if debug: # Determine save paths based on prefix (save to visualization folder) if save_prefix: final_result_path = os.path.join(viz_folder, f"{save_prefix}_final_result.png") composite_result_path = os.path.join(viz_folder, f"{save_prefix}_composite_result.png") else: final_result_path = os.path.join(viz_folder, "final_result.png") composite_result_path = os.path.join(viz_folder, "composite_result.png") # Visualize final result (only for 2 nodes, skip for N nodes) # TODO: Extend visualize_final_result to support N nodes if num_nodes == 2: mask1_orig = load_binary_mask_from_rgba(png_list[0]) mask2_orig = load_binary_mask_from_rgba(png_list[1]) if mask1_orig is not None and mask2_orig is not None: visualize_final_result(softmasks_f[0], softmasks_f[1], distances_f[0], distances_f[1], final_bboxes[0], final_bboxes[1], mask1_orig, mask2_orig, Wc, Hc, save_path=final_result_path) # Save composite image with original images # Use original paths if provided, otherwise use the paths passed to optimize save_png_list = [] for i in range(num_nodes): orig_png = original_png_list[i] if i < len(original_png_list) and original_png_list[i] is not None else png_list[i] save_png_list.append(orig_png) if all(png is not None for png in save_png_list): save_composite_image(png_list=save_png_list, bbox_list=final_bboxes, Wc=Wc, Hc=Hc, save_path=composite_result_path) else: print(f"Warning: Skipping composite image save - some image paths are None") return final_bboxes