Spaces:
Running
Running
File size: 7,735 Bytes
c8a046c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 |
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 |