import cv2 import numpy as np from typing import List, Dict, Tuple import yaml class DamageVisualizer: """Visualize detection and comparison results""" def __init__(self, config_path: str = "config.yaml"): """Initialize visualizer with configuration""" with open(config_path, 'r') as f: self.config = yaml.safe_load(f) self.line_thickness = self.config['visualization']['line_thickness'] self.font_scale = self.config['visualization']['font_scale'] self.colors = self.config['visualization']['colors'] def draw_detections(self, image: np.ndarray, detections: Dict, color_type: str = 'new_damage') -> np.ndarray: """ Draw bounding boxes and labels on image Args: image: Input image detections: Detection results color_type: Type of color to use ('new_damage', 'existing_damage', 'matched_damage') Returns: Image with drawn detections """ img_copy = image.copy() color = self.colors.get(color_type, [255, 0, 0]) for i, box in enumerate(detections['boxes']): x1, y1, x2, y2 = box label = f"{detections['classes'][i]} ({detections['confidences'][i]:.2f})" # Draw rectangle cv2.rectangle(img_copy, (x1, y1), (x2, y2), color, self.line_thickness) # Draw label background label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, self.font_scale, self.line_thickness) cv2.rectangle(img_copy, (x1, y1 - label_size[1] - 5), (x1 + label_size[0], y1), color, -1) # Draw label text cv2.putText(img_copy, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, self.font_scale, [255, 255, 255], self.line_thickness) return img_copy def create_comparison_visualization(self, before_img: np.ndarray, after_img: np.ndarray, before_detections: Dict, after_detections: Dict, comparison_result: Dict) -> np.ndarray: """ Create side-by-side comparison visualization Args: before_img, after_img: Input images before_detections, after_detections: Detection results comparison_result: Comparison analysis results Returns: Combined visualization image """ # Draw matched damages in yellow before_vis = before_img.copy() after_vis = after_img.copy() # Draw matched damages for match in comparison_result['matched_damages']: # Draw on before image x1, y1, x2, y2 = match['box_before'] cv2.rectangle(before_vis, (x1, y1), (x2, y2), self.colors['matched_damage'], self.line_thickness) # Draw on after image x1, y1, x2, y2 = match['box_after'] cv2.rectangle(after_vis, (x1, y1), (x2, y2), self.colors['matched_damage'], self.line_thickness) # Draw repaired damages (only on before) in green for damage in comparison_result['repaired_damages']: x1, y1, x2, y2 = damage['box'] cv2.rectangle(before_vis, (x1, y1), (x2, y2), self.colors['existing_damage'], self.line_thickness) cv2.putText(before_vis, "REPAIRED", (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, self.font_scale, self.colors['existing_damage'], self.line_thickness) # Draw new damages (only on after) in red for damage in comparison_result['new_damages']: x1, y1, x2, y2 = damage['box'] cv2.rectangle(after_vis, (x1, y1), (x2, y2), self.colors['new_damage'], self.line_thickness + 1) cv2.putText(after_vis, "NEW!", (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, self.font_scale * 1.5, self.colors['new_damage'], self.line_thickness) # Combine images side by side h1, w1 = before_vis.shape[:2] h2, w2 = after_vis.shape[:2] max_height = max(h1, h2) # Resize if needed if h1 != max_height: before_vis = cv2.resize(before_vis, (int(w1 * max_height / h1), max_height)) if h2 != max_height: after_vis = cv2.resize(after_vis, (int(w2 * max_height / h2), max_height)) # Create combined image combined = np.hstack([before_vis, after_vis]) # Add status text status_height = 100 status_img = np.ones((status_height, combined.shape[1], 3), dtype=np.uint8) * 255 # Add case message case_color = (0, 128, 0) if 'SUCCESS' in comparison_result['case'] else (0, 0, 255) cv2.putText(status_img, comparison_result['message'], (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.8, case_color, 2) # Add statistics stats_text = f"Before: {comparison_result['statistics']['total_before']} | " \ f"After: {comparison_result['statistics']['total_after']} | " \ f"Matched: {comparison_result['statistics']['matched']} | " \ f"New: {comparison_result['statistics']['new']}" cv2.putText(status_img, stats_text, (20, 70), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1) # Combine with status final_image = np.vstack([status_img, combined]) return final_image def create_summary_grid(self, comparison_results: List[Dict], image_pairs: List[Tuple[np.ndarray, np.ndarray]]) -> np.ndarray: """ Create a grid visualization of all 6 position comparisons Args: comparison_results: List of comparison results for each position image_pairs: List of (before, after) image pairs Returns: Grid visualization of all positions """ grid_images = [] for i, (result, (before_img, after_img)) in enumerate(zip(comparison_results, image_pairs)): # Create mini comparison for each position target_size = (300, 200) # Smaller size for grid before_small = cv2.resize(before_img, target_size) after_small = cv2.resize(after_img, target_size) # Add position label position_label = f"Position {i+1}" cv2.putText(before_small, position_label, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) cv2.putText(after_small, position_label, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) # Add case indicator case_color = (0, 255, 0) if 'SUCCESS' in result['case'] else (0, 0, 255) if 'NEW_DAMAGE' in result['case']: case_color = (0, 0, 255) cv2.rectangle(after_small, (0, 0), (target_size[0], 30), case_color, -1) cv2.putText(after_small, result['case'][:10], (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1) pair_img = np.hstack([before_small, after_small]) grid_images.append(pair_img) # Create 2x3 grid row1 = np.hstack(grid_images[:3]) row2 = np.hstack(grid_images[3:]) grid = np.vstack([row1, row2]) return grid