import matplotlib matplotlib.use('Agg') # Non-interactive backend for web apps import matplotlib.pyplot as plt from matplotlib.patches import FancyBboxPatch, FancyArrowPatch import xml.etree.ElementTree as ET import cv2 import re from typing import Dict, Set, Tuple # ========================================== # 1. Tree Visualizer (Refactored from your upload) # ========================================== class XMLTreeVisualizer: # Single color scheme for all nodes DEFAULT_COLOR = {'fill': '#E3F2FD', 'border': '#1976D2', 'text': '#000000'} HIGHLIGHT_COLOR = {'fill': '#FFF59D', 'border': '#F57F17', 'text': '#000000'} # Yellow for active nodes def visualize(self, xml_path: str, output_path: str, visible_text: Set[str] = None, active_elements: Set = None): """Generates tree visualization. If visible_text and active_elements are provided, highlights active nodes.""" tree = ET.parse(xml_path) root = tree.getroot() # Calculate Layout positions = self._calculate_layout(root) max_depth = max(p['level'] for p in positions.values()) # Setup Figure with larger size for better text readability fig, ax = plt.subplots(figsize=(24, 18)) ax.set_xlim(-len(positions)*0.5, len(positions)*0.5) ax.set_ylim(-max_depth * 2 - 2, 2) ax.axis('off') # Draw Edges self._draw_edges(ax, positions, root) # Draw Nodes self._draw_nodes(ax, positions, visible_text, active_elements) plt.title("XML Tree Structure" + (" (Active Nodes Highlighted)" if active_elements else ""), fontsize=16) plt.tight_layout() plt.savefig(output_path, dpi=150, bbox_inches='tight') plt.close(fig) def _calculate_layout(self, root, x=0, y=0, level=0, spacing=2.5): positions = {} def _get_width(node): children = list(node) if not children: return 1.0 return sum(_get_width(c) for c in children) def _assign_pos(node, curr_x, curr_y, curr_level): positions[id(node)] = {'x': curr_x, 'y': curr_y, 'level': curr_level, 'element': node} children = list(node) if not children: return width = sum(_get_width(c) for c in children) start_x = curr_x - (width * spacing / 2) current_offset = 0 for child in children: child_w = _get_width(child) child_x = start_x + (current_offset + child_w/2) * spacing _assign_pos(child, child_x, curr_y - 2, curr_level + 1) current_offset += child_w _assign_pos(root, x, y, level) return positions def _draw_edges(self, ax, positions, node): node_id = id(node) if node_id not in positions: return parent_pos = positions[node_id] for child in node: child_id = id(child) if child_id in positions: child_pos = positions[child_id] arrow = FancyArrowPatch( (parent_pos['x'], parent_pos['y']), (child_pos['x'], child_pos['y']), arrowstyle='-', color='#555', linewidth=1, zorder=1 ) ax.add_patch(arrow) self._draw_edges(ax, positions, child) def _draw_nodes(self, ax, positions, visible_text, active_elements): for pid, info in positions.items(): elem = info['element'] text = elem.get('text', '').strip() # Highlight Logic: Check if this element is in the active_elements set is_highlight = False if active_elements and elem in active_elements: is_highlight = True # Use single color scheme if is_highlight: face = self.HIGHLIGHT_COLOR['fill'] edge = self.HIGHLIGHT_COLOR['border'] lw = 3 else: face = self.DEFAULT_COLOR['fill'] edge = self.DEFAULT_COLOR['border'] lw = 1 # Calculate box size based on text length if text: # Use actual text, wrap if too long display_text = text # Wrap text if longer than 20 characters if len(display_text) > 20: # Try to break at word boundaries words = display_text.split() lines = [] current_line = "" for word in words: if len(current_line + " " + word) <= 20: current_line = current_line + " " + word if current_line else word else: if current_line: lines.append(current_line) current_line = word if current_line: lines.append(current_line) display_text = "\n".join(lines[:2]) # Max 2 lines # Calculate box dimensions based on text num_lines = display_text.count('\n') + 1 text_width = max(len(line) for line in display_text.split('\n')) if display_text else 0 box_width = max(1.2, min(3.0, text_width * 0.15)) box_height = max(0.8, 0.6 + (num_lines - 1) * 0.4) else: display_text = "" box_width = 1.2 box_height = 0.8 # Draw Box box = FancyBboxPatch( (info['x']-box_width/2, info['y']-box_height/2), box_width, box_height, boxstyle="round,pad=0.1", facecolor=face, edgecolor=edge, linewidth=lw, zorder=2 ) ax.add_patch(box) # Draw Text Label - only show actual text from XML, make it readable if display_text: # Use larger, readable font - adjust based on text length max_line_len = max(len(line) for line in display_text.split('\n')) if '\n' in display_text else len(display_text) if max_line_len <= 10: fontsize = 11 elif max_line_len <= 15: fontsize = 10 else: fontsize = 9 ax.text(info['x'], info['y'], display_text, ha='center', va='center', fontsize=fontsize, zorder=3, family='sans-serif', weight='normal') # ========================================== # 2. Bounding Box Visualizer (Refactored) # ========================================== class BoundingBoxVisualizer: COLORS = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0)] def visualize(self, image_path: str, xml_path: str, output_path: str): img = cv2.imread(image_path) if img is None: return tree = ET.parse(xml_path) idx = 0 for elem in tree.getroot().iter(): bounds = elem.get('bounds') if bounds: # Parse [x1,y1][x2,y2] match = re.match(r'\[(\d+),(\d+)\]\[(\d+),(\d+)\]', bounds) if match: x1, y1, x2, y2 = map(int, match.groups()) color = self.COLORS[idx % len(self.COLORS)] cv2.rectangle(img, (x1, y1), (x2, y2), color, 2) idx += 1 cv2.imwrite(output_path, img)