src/visualization.py

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