254 lines
7.6 KiB
Python
254 lines
7.6 KiB
Python
"""
|
|
Debug visualization for detections.
|
|
"""
|
|
|
|
import cv2
|
|
import logging
|
|
import numpy as np
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from typing import List, Any, Tuple, Optional, Dict
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DebugVisualizer:
|
|
"""
|
|
Visualize detections with bounding boxes and labels.
|
|
|
|
Features:
|
|
- Draw bounding boxes with class labels
|
|
- Confidence scores
|
|
- FPS counter
|
|
- Save debug frames
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
output_dir: str = "output/debug",
|
|
camera_name: str = "default",
|
|
draw_boxes: bool = True,
|
|
draw_labels: bool = True,
|
|
draw_confidence: bool = True,
|
|
box_color: Tuple[int, int, int] = (0, 255, 0),
|
|
box_thickness: int = 2,
|
|
font_scale: float = 0.5,
|
|
save_interval: int = 10,
|
|
):
|
|
"""
|
|
Initialize debug visualizer.
|
|
|
|
Args:
|
|
output_dir: Directory to save debug frames
|
|
camera_name: Camera name
|
|
draw_boxes: Draw bounding boxes
|
|
draw_labels: Draw class labels
|
|
draw_confidence: Draw confidence scores
|
|
box_color: Default box color (BGR)
|
|
box_thickness: Box line thickness
|
|
font_scale: Font scale for labels
|
|
save_interval: Save every N frames (0 = save all)
|
|
"""
|
|
self.output_dir = Path(output_dir) / camera_name
|
|
self.camera_name = camera_name
|
|
self.draw_boxes = draw_boxes
|
|
self.draw_labels = draw_labels
|
|
self.draw_confidence = draw_confidence
|
|
self.box_color = box_color
|
|
self.box_thickness = box_thickness
|
|
self.font_scale = font_scale
|
|
self.save_interval = save_interval
|
|
|
|
self.output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
self._frame_count = 0
|
|
self._class_colors: Dict[int, Tuple[int, int, int]] = {}
|
|
|
|
def _get_class_color(self, class_id: int) -> Tuple[int, int, int]:
|
|
"""Get consistent color for class ID."""
|
|
if class_id not in self._class_colors:
|
|
# Generate deterministic color from class ID
|
|
np.random.seed(class_id + 42)
|
|
self._class_colors[class_id] = tuple(
|
|
int(x) for x in np.random.randint(50, 255, 3)
|
|
)
|
|
return self._class_colors[class_id]
|
|
|
|
def draw_detections(
|
|
self,
|
|
frame: np.ndarray,
|
|
detections: List[Any],
|
|
use_class_colors: bool = True,
|
|
) -> np.ndarray:
|
|
"""
|
|
Draw detection boxes on frame.
|
|
|
|
Args:
|
|
frame: Input frame (BGR)
|
|
detections: List of Detection objects
|
|
use_class_colors: Use different colors per class
|
|
|
|
Returns:
|
|
Frame with annotations
|
|
"""
|
|
output = frame.copy()
|
|
|
|
for det in detections:
|
|
# Extract detection info
|
|
if hasattr(det, 'bbox'):
|
|
x1, y1 = int(det.bbox.x1), int(det.bbox.y1)
|
|
x2, y2 = int(det.bbox.x2), int(det.bbox.y2)
|
|
class_id = det.class_id
|
|
class_name = det.class_name
|
|
conf = det.confidence
|
|
else:
|
|
bbox = det.get('bbox', [0, 0, 0, 0])
|
|
x1, y1, x2, y2 = [int(x) for x in bbox]
|
|
class_id = det.get('class_id', 0)
|
|
class_name = det.get('class_name', str(class_id))
|
|
conf = det.get('confidence', 0)
|
|
|
|
# Get color
|
|
if use_class_colors:
|
|
color = self._get_class_color(class_id)
|
|
else:
|
|
color = self.box_color
|
|
|
|
# Draw box
|
|
if self.draw_boxes:
|
|
cv2.rectangle(output, (x1, y1), (x2, y2), color, self.box_thickness)
|
|
|
|
# Build label
|
|
if self.draw_labels or self.draw_confidence:
|
|
label_parts = []
|
|
if self.draw_labels:
|
|
label_parts.append(class_name)
|
|
if self.draw_confidence:
|
|
label_parts.append(f"{conf:.2f}")
|
|
|
|
label = " ".join(label_parts)
|
|
|
|
# Calculate label size
|
|
(label_w, label_h), baseline = cv2.getTextSize(
|
|
label, cv2.FONT_HERSHEY_SIMPLEX, self.font_scale, 1
|
|
)
|
|
|
|
# Draw label background
|
|
cv2.rectangle(
|
|
output,
|
|
(x1, y1 - label_h - 6),
|
|
(x1 + label_w + 4, y1),
|
|
color,
|
|
-1
|
|
)
|
|
|
|
# Draw label text
|
|
cv2.putText(
|
|
output,
|
|
label,
|
|
(x1 + 2, y1 - 4),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
self.font_scale,
|
|
(255, 255, 255),
|
|
1,
|
|
cv2.LINE_AA,
|
|
)
|
|
|
|
return output
|
|
|
|
def draw_info(
|
|
self,
|
|
frame: np.ndarray,
|
|
fps: Optional[float] = None,
|
|
frame_id: Optional[int] = None,
|
|
detection_count: Optional[int] = None,
|
|
) -> np.ndarray:
|
|
"""
|
|
Draw info overlay on frame.
|
|
|
|
Args:
|
|
frame: Input frame
|
|
fps: FPS value
|
|
frame_id: Frame index
|
|
detection_count: Number of detections
|
|
|
|
Returns:
|
|
Frame with info overlay
|
|
"""
|
|
output = frame.copy()
|
|
y = 25
|
|
|
|
if fps is not None:
|
|
text = f"FPS: {fps:.1f}"
|
|
cv2.putText(output, text, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
|
|
y += 25
|
|
|
|
if frame_id is not None:
|
|
text = f"Frame: {frame_id}"
|
|
cv2.putText(output, text, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
|
|
y += 25
|
|
|
|
if detection_count is not None:
|
|
text = f"Objects: {detection_count}"
|
|
cv2.putText(output, text, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
|
|
|
|
return output
|
|
|
|
def process_frame(
|
|
self,
|
|
frame: np.ndarray,
|
|
detections: List[Any],
|
|
frame_id: int = 0,
|
|
fps: Optional[float] = None,
|
|
save: bool = True,
|
|
) -> np.ndarray:
|
|
"""
|
|
Process frame with all visualizations.
|
|
|
|
Args:
|
|
frame: Input frame
|
|
detections: List of detections
|
|
frame_id: Frame index
|
|
fps: Current FPS
|
|
save: Save to file
|
|
|
|
Returns:
|
|
Visualized frame
|
|
"""
|
|
# Draw detections
|
|
output = self.draw_detections(frame, detections)
|
|
|
|
# Draw info overlay
|
|
output = self.draw_info(output, fps, frame_id, len(detections))
|
|
|
|
# Save if needed
|
|
self._frame_count += 1
|
|
if save and self.save_interval > 0:
|
|
if self._frame_count % self.save_interval == 0:
|
|
self.save_frame(output, frame_id)
|
|
|
|
return output
|
|
|
|
def save_frame(
|
|
self,
|
|
frame: np.ndarray,
|
|
frame_id: int = 0,
|
|
) -> str:
|
|
"""
|
|
Save debug frame to file.
|
|
|
|
Args:
|
|
frame: Frame to save
|
|
frame_id: Frame index
|
|
|
|
Returns:
|
|
Path to saved file
|
|
"""
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
filename = f"debug_{timestamp}_{frame_id:06d}.jpg"
|
|
filepath = self.output_dir / filename
|
|
|
|
cv2.imwrite(str(filepath), frame, [cv2.IMWRITE_JPEG_QUALITY, 85])
|
|
|
|
return str(filepath)
|