415 lines
15 KiB
Python
415 lines
15 KiB
Python
import cv2
|
|
import numpy as np
|
|
from typing import Dict, List, Tuple, Any, Optional
|
|
from PySide6.QtGui import QImage, QPixmap
|
|
from PySide6.QtCore import Qt
|
|
|
|
# Color mapping for traffic-related classes
|
|
COLORS = {
|
|
'person': (255, 165, 0), # Orange
|
|
'bicycle': (255, 0, 255), # Magenta
|
|
'car': (0, 255, 0), # Green
|
|
'motorcycle': (255, 255, 0), # Cyan
|
|
'bus': (0, 0, 255), # Red
|
|
'truck': (0, 128, 255), # Orange-Blue
|
|
'traffic light': (0, 165, 255), # Orange
|
|
'stop sign': (0, 0, 139), # Dark Red
|
|
'parking meter': (128, 0, 128), # Purple
|
|
'default': (0, 255, 255) # Yellow as default
|
|
}
|
|
|
|
# Enhanced class colors for consistent visualization
|
|
def get_enhanced_class_color(class_name: str, class_id: int = -1) -> Tuple[int, int, int]:
|
|
"""
|
|
Get color for class with enhanced mapping (traffic classes only)
|
|
|
|
Args:
|
|
class_name: Name of the detected class
|
|
class_id: COCO class ID
|
|
|
|
Returns:
|
|
BGR color tuple
|
|
"""
|
|
# Only traffic class IDs/colors
|
|
enhanced_colors = {
|
|
0: (255, 165, 0), # person - Orange
|
|
1: (255, 0, 255), # bicycle - Magenta
|
|
2: (0, 255, 0), # car - Green
|
|
3: (255, 255, 0), # motorcycle - Cyan
|
|
4: (0, 0, 255), # bus - Red
|
|
5: (0, 128, 255), # truck - Orange-Blue
|
|
6: (0, 165, 255), # traffic light - Orange
|
|
7: (0, 0, 139), # stop sign - Dark Red
|
|
8: (128, 0, 128), # parking meter - Purple
|
|
}
|
|
|
|
# Get color from class name if available
|
|
if class_name and class_name.lower() in COLORS:
|
|
return COLORS[class_name.lower()]
|
|
|
|
# Get color from class ID if available
|
|
if isinstance(class_id, int) and class_id in enhanced_colors:
|
|
return enhanced_colors[class_id]
|
|
|
|
# Default color
|
|
return COLORS['default']
|
|
|
|
def enhanced_draw_detections(frame: np.ndarray, detections: List[Dict],
|
|
draw_labels: bool = True, draw_confidence: bool = True) -> np.ndarray:
|
|
"""
|
|
Enhanced version of draw_detections with better visualization
|
|
|
|
Args:
|
|
frame: Input video frame
|
|
detections: List of detection dictionaries
|
|
draw_labels: Whether to draw class labels
|
|
draw_confidence: Whether to draw confidence scores
|
|
|
|
Returns:
|
|
Annotated frame
|
|
"""
|
|
if frame is None or not isinstance(frame, np.ndarray) or frame.size == 0:
|
|
print("Warning: Invalid frame provided to enhanced_draw_detections")
|
|
return np.zeros((300, 300, 3), dtype=np.uint8) # Return blank frame as fallback
|
|
|
|
annotated_frame = frame.copy()
|
|
|
|
# Handle case when detections is None or empty
|
|
if detections is None or len(detections) == 0:
|
|
return annotated_frame
|
|
|
|
# Get frame dimensions for validation
|
|
h, w = frame.shape[:2]
|
|
|
|
for detection in detections:
|
|
if not isinstance(detection, dict):
|
|
continue
|
|
|
|
try:
|
|
# Skip detection if it doesn't have bbox or has invalid confidence
|
|
if 'bbox' not in detection:
|
|
continue
|
|
|
|
# Skip if confidence is below threshold (don't rely on external filtering)
|
|
confidence = detection.get('confidence', 0.0)
|
|
if confidence < 0.1: # Apply a minimal threshold to ensure we're not drawing noise
|
|
continue
|
|
|
|
bbox = detection['bbox']
|
|
class_name = detection.get('class_name', 'unknown')
|
|
class_id = detection.get('class_id', -1)
|
|
|
|
# Get color for class
|
|
color = get_enhanced_class_color(class_name, class_id)
|
|
|
|
# Ensure bbox has enough coordinates and they are numeric values
|
|
if len(bbox) < 4 or not all(isinstance(coord, (int, float)) for coord in bbox[:4]):
|
|
continue
|
|
|
|
# Convert coordinates to integers
|
|
try:
|
|
x1, y1, x2, y2 = map(int, bbox[:4])
|
|
except (ValueError, TypeError):
|
|
print(f"Warning: Invalid bbox format: {bbox}")
|
|
continue
|
|
|
|
# Validate coordinates are within frame bounds
|
|
x1 = max(0, min(x1, w-1))
|
|
y1 = max(0, min(y1, h-1))
|
|
x2 = max(0, min(x2, w))
|
|
y2 = max(0, min(y2, h))
|
|
|
|
# Ensure x2 > x1 and y2 > y1 (at least 1 pixel width/height)
|
|
if x2 <= x1 or y2 <= y1:
|
|
# Instead of skipping, fix the coordinates to ensure at least 1 pixel width/height
|
|
x2 = max(x1 + 1, x2)
|
|
y2 = max(y1 + 1, y2)
|
|
|
|
# Draw bounding box with thicker line for better visibility
|
|
cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), color, 2)
|
|
|
|
# Prepare label text
|
|
label_parts = []
|
|
if draw_labels:
|
|
# Display proper class name
|
|
display_name = class_name.replace('_', ' ').title()
|
|
label_parts.append(display_name)
|
|
|
|
# Add tracking ID if available
|
|
track_id = detection.get('track_id')
|
|
if track_id is not None:
|
|
label_parts[-1] += f" #{track_id}"
|
|
|
|
if draw_confidence and confidence > 0:
|
|
label_parts.append(f"{confidence:.2f}")
|
|
|
|
# Draw traffic light color indicator if available
|
|
if class_name == 'traffic light' and 'traffic_light_color' in detection:
|
|
light_color = detection['traffic_light_color']
|
|
|
|
# Add traffic light color to label
|
|
if light_color != 'unknown':
|
|
# Set color indicator based on traffic light state
|
|
if light_color == 'red':
|
|
color_indicator = (0, 0, 255) # Red
|
|
label_parts.append("🔴 RED")
|
|
elif light_color == 'yellow':
|
|
color_indicator = (0, 255, 255) # Yellow
|
|
label_parts.append("🟡 YELLOW")
|
|
elif light_color == 'green':
|
|
color_indicator = (0, 255, 0) # Green
|
|
label_parts.append("🟢 GREEN")
|
|
|
|
# Draw traffic light visual indicator (circle with detected color)
|
|
circle_y = y1 - 15
|
|
circle_x = x1 + 10
|
|
circle_radius = 10
|
|
|
|
if light_color == 'red':
|
|
cv2.circle(annotated_frame, (circle_x, circle_y), circle_radius, (0, 0, 255), -1)
|
|
elif light_color == 'yellow':
|
|
cv2.circle(annotated_frame, (circle_x, circle_y), circle_radius, (0, 255, 255), -1)
|
|
elif light_color == 'green':
|
|
cv2.circle(annotated_frame, (circle_x, circle_y), circle_radius, (0, 255, 0), -1)
|
|
|
|
# Draw label if we have any text
|
|
if label_parts:
|
|
label = " ".join(label_parts)
|
|
|
|
try:
|
|
# Get text size for background
|
|
(text_width, text_height), baseline = cv2.getTextSize(
|
|
label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2
|
|
)
|
|
|
|
# Ensure label position is within frame
|
|
text_y = max(text_height + 10, y1)
|
|
|
|
# Draw label background (use colored background)
|
|
bg_color = tuple(int(c * 0.7) for c in color) # Darker version of box color
|
|
cv2.rectangle(
|
|
annotated_frame,
|
|
(x1, text_y - text_height - 10),
|
|
(x1 + text_width + 10, text_y),
|
|
bg_color,
|
|
-1
|
|
)
|
|
# Draw label text (white text on colored background)
|
|
cv2.putText(
|
|
annotated_frame,
|
|
label,
|
|
(x1 + 5, text_y - 5),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.6,
|
|
(255, 255, 255), # White text
|
|
2
|
|
)
|
|
except Exception as e:
|
|
print(f"Error drawing label: {e}")
|
|
|
|
except Exception as e:
|
|
print(f"Error drawing detection: {e}")
|
|
continue
|
|
|
|
return annotated_frame
|
|
|
|
def draw_performance_overlay(frame: np.ndarray, metrics: Dict) -> np.ndarray:
|
|
"""
|
|
Draw enhanced performance metrics overlay on the frame.
|
|
|
|
Args:
|
|
frame: Input video frame
|
|
metrics: Dictionary of performance metrics
|
|
|
|
Returns:
|
|
Annotated frame
|
|
"""
|
|
if frame is None or not isinstance(frame, np.ndarray):
|
|
return np.zeros((300, 300, 3), dtype=np.uint8)
|
|
|
|
annotated_frame = frame.copy()
|
|
height, width = annotated_frame.shape[:2]
|
|
|
|
# Create semi-transparent overlay for metrics panel
|
|
overlay = annotated_frame.copy()
|
|
|
|
# Calculate panel size based on metrics count
|
|
text_height = 25
|
|
padding = 10
|
|
metrics_count = len(metrics)
|
|
panel_height = metrics_count * text_height + 2 * padding
|
|
panel_width = 220 # Fixed width
|
|
|
|
# Position panel at bottom left
|
|
panel_x = 10
|
|
panel_y = height - panel_height - 10
|
|
|
|
# Draw background panel with transparency
|
|
cv2.rectangle(
|
|
overlay,
|
|
(panel_x, panel_y),
|
|
(panel_x + panel_width, panel_y + panel_height),
|
|
(0, 0, 0),
|
|
-1
|
|
)
|
|
|
|
# Apply transparency
|
|
alpha = 0.7
|
|
cv2.addWeighted(overlay, alpha, annotated_frame, 1 - alpha, 0, annotated_frame)
|
|
|
|
# Draw metrics with custom formatting
|
|
text_y = panel_y + padding + text_height
|
|
for metric, value in metrics.items():
|
|
# Format metric name and value
|
|
metric_text = f"{metric}: {value}"
|
|
|
|
# Choose color based on metric type
|
|
if "FPS" in metric:
|
|
color = (0, 255, 0) # Green for FPS
|
|
elif "ms" in str(value):
|
|
color = (0, 255, 255) # Yellow for timing metrics
|
|
else:
|
|
color = (255, 255, 255) # White for other metrics
|
|
|
|
# Draw text with drop shadow for better readability
|
|
cv2.putText(
|
|
annotated_frame,
|
|
metric_text,
|
|
(panel_x + padding + 1, text_y + 1),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.6,
|
|
(0, 0, 0), # Black shadow
|
|
2
|
|
)
|
|
cv2.putText(
|
|
annotated_frame,
|
|
metric_text,
|
|
(panel_x + padding, text_y),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.6,
|
|
color,
|
|
2
|
|
)
|
|
text_y += text_height
|
|
|
|
return annotated_frame
|
|
|
|
def resize_frame_for_display(frame: np.ndarray, max_width: int = 1280, max_height: int = 720) -> np.ndarray:
|
|
"""
|
|
Resize frame for display while maintaining aspect ratio.
|
|
|
|
Args:
|
|
frame: Input video frame
|
|
max_width: Maximum display width
|
|
max_height: Maximum display height
|
|
|
|
Returns:
|
|
Resized frame
|
|
"""
|
|
if frame is None:
|
|
return np.zeros((300, 300, 3), dtype=np.uint8)
|
|
|
|
height, width = frame.shape[:2]
|
|
|
|
# No resize needed if image is already smaller than max dimensions
|
|
if width <= max_width and height <= max_height:
|
|
return frame
|
|
|
|
# Calculate scale factor to fit within max dimensions
|
|
scale_width = max_width / width if width > max_width else 1.0
|
|
scale_height = max_height / height if height > max_height else 1.0
|
|
|
|
# Use the smaller scale to ensure image fits within bounds
|
|
scale = min(scale_width, scale_height)
|
|
|
|
# Resize using calculated scale
|
|
new_width = int(width * scale)
|
|
new_height = int(height * scale)
|
|
|
|
return cv2.resize(frame, (new_width, new_height), interpolation=cv2.INTER_AREA)
|
|
|
|
def enhanced_cv_to_qimage(cv_img: np.ndarray) -> QImage:
|
|
"""
|
|
Enhanced converter from OpenCV image to QImage with robust error handling.
|
|
|
|
Args:
|
|
cv_img: OpenCV image (numpy array)
|
|
|
|
Returns:
|
|
QImage object
|
|
"""
|
|
if cv_img is None or not isinstance(cv_img, np.ndarray):
|
|
print("Warning: Invalid image in enhanced_cv_to_qimage")
|
|
# Return a small black image as fallback
|
|
return QImage(10, 10, QImage.Format_RGB888)
|
|
|
|
try:
|
|
# Get image dimensions and verify its validity
|
|
h, w, ch = cv_img.shape
|
|
if h <= 0 or w <= 0 or ch != 3:
|
|
raise ValueError(f"Invalid image dimensions: {h}x{w}x{ch}")
|
|
|
|
# OpenCV uses BGR, Qt uses RGB format, so convert
|
|
rgb_image = cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB)
|
|
|
|
# Calculate bytes per line
|
|
bytes_per_line = ch * w
|
|
|
|
# Use numpy array data directly
|
|
# This avoids a copy, but ensures the data is properly aligned
|
|
# by creating a contiguous array
|
|
contiguous_data = np.ascontiguousarray(rgb_image)
|
|
|
|
# Create QImage from numpy array
|
|
q_image = QImage(contiguous_data.data, w, h, bytes_per_line, QImage.Format_RGB888)
|
|
|
|
# Create a copy to ensure the data stays valid when returning
|
|
return q_image.copy()
|
|
except Exception as e:
|
|
print(f"Error in enhanced_cv_to_qimage: {e}")
|
|
# Return a small black image as fallback
|
|
return QImage(10, 10, QImage.Format_RGB888)
|
|
|
|
def enhanced_cv_to_pixmap(cv_img: np.ndarray, target_width: int = None) -> QPixmap:
|
|
"""
|
|
Enhanced converter from OpenCV image to QPixmap with robust error handling.
|
|
|
|
Args:
|
|
cv_img: OpenCV image (numpy array)
|
|
target_width: Optional width to resize to (maintains aspect ratio)
|
|
|
|
Returns:
|
|
QPixmap object
|
|
"""
|
|
if cv_img is None or not isinstance(cv_img, np.ndarray):
|
|
print("Warning: Invalid image in enhanced_cv_to_pixmap")
|
|
# Create an empty pixmap with visual indication of error
|
|
empty_pixmap = QPixmap(640, 480)
|
|
empty_pixmap.fill(Qt.black)
|
|
return empty_pixmap
|
|
|
|
try:
|
|
# First convert to QImage
|
|
q_image = enhanced_cv_to_qimage(cv_img)
|
|
|
|
if q_image.isNull():
|
|
raise ValueError("Generated null QImage")
|
|
|
|
# Resize if needed
|
|
if target_width and q_image.width() > target_width:
|
|
q_image = q_image.scaledToWidth(target_width, Qt.SmoothTransformation)
|
|
|
|
# Convert to QPixmap
|
|
pixmap = QPixmap.fromImage(q_image)
|
|
|
|
if pixmap.isNull():
|
|
raise ValueError("Generated null QPixmap")
|
|
|
|
return pixmap
|
|
except Exception as e:
|
|
print(f"Error in enhanced_cv_to_pixmap: {e}")
|
|
# Create an empty pixmap with visual indication of error
|
|
empty_pixmap = QPixmap(640, 480)
|
|
empty_pixmap.fill(Qt.black)
|
|
return empty_pixmap
|