Traffic-Intersection-Monitoring/qt_app_pyside1/ui/widgets/detection_overlay_widget.py

"""
Detection Overlay Widget - Displays object detection bounding boxes and labels
"""

from PySide6.QtWidgets import QWidget, QLabel, QVBoxLayout, QHBoxLayout, QFrame
from PySide6.QtCore import Qt, Signal, QRect, QPoint, QTimer
from PySide6.QtGui import QPainter, QPen, QBrush, QColor, QFont
import json

class DetectionOverlayWidget(QWidget):
    """
    Widget for displaying object detection overlays on video frames
    
    Features:
    - Bounding box visualization
    - Class labels and confidence scores
    - Tracking IDs for multi-object tracking
    - Color-coded detection classes
    - Customizable overlay styles
    """
    
    # Signals
    detection_clicked = Signal(dict)  # Emitted when a detection is clicked
    
    def __init__(self, parent=None):
        super().__init__(parent)
        
        self.detections = []
        self.class_colors = {
            'person': QColor(255, 0, 0),      # Red
            'bicycle': QColor(0, 255, 0),     # Green
            'car': QColor(0, 0, 255),         # Blue
            'motorcycle': QColor(255, 255, 0), # Yellow
            'bus': QColor(255, 0, 255),       # Magenta
            'truck': QColor(0, 255, 255),     # Cyan
            'traffic_light': QColor(255, 165, 0), # Orange
            'stop_sign': QColor(128, 0, 128),  # Purple
        }
        
        self.show_labels = True
        self.show_confidence = True
        self.show_tracking_id = True
        self.overlay_opacity = 0.8
        
        # Make widget transparent for overlay
        self.setAttribute(Qt.WA_TransparentForMouseEvents, False)
        self.setStyleSheet("background-color: transparent;")
    
    def set_detections(self, detections):
        """
        Set detection results to display
        
        Args:
            detections: List of detection dictionaries with format:
                {
                    'bbox': [x1, y1, x2, y2],
                    'class': 'person',
                    'confidence': 0.95,
                    'track_id': 123 (optional)
                }
        """
        self.detections = detections
        self.update()  # Trigger repaint
    
    def set_overlay_options(self, show_labels=True, show_confidence=True, 
                          show_tracking_id=True, opacity=0.8):
        """Configure overlay display options"""
        self.show_labels = show_labels
        self.show_confidence = show_confidence
        self.show_tracking_id = show_tracking_id
        self.overlay_opacity = opacity
        self.update()
    
    def add_class_color(self, class_name, color):
        """Add or update color for a detection class"""
        self.class_colors[class_name] = color
    
    def paintEvent(self, event):
        """Paint detection overlays"""
        if not self.detections:
            return
        
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)
        
        # Set up font for labels
        font = QFont("Arial", 10, QFont.Bold)
        painter.setFont(font)
        
        for detection in self.detections:
            self._draw_detection(painter, detection)
    
    def _draw_detection(self, painter, detection):
        """Draw a single detection"""
        # Extract detection info
        bbox = detection.get('bbox', [0, 0, 100, 100])
        class_name = detection.get('class', 'unknown')
        confidence = detection.get('confidence', 0.0)
        track_id = detection.get('track_id', None)
        
        x1, y1, x2, y2 = bbox
        
        # Get color for this class
        color = self.class_colors.get(class_name, QColor(128, 128, 128))
        
        # Set up pen for bounding box
        pen = QPen(color, 3)
        painter.setPen(pen)
        
        # Draw bounding box
        rect = QRect(int(x1), int(y1), int(x2-x1), int(y2-y1))
        painter.drawRect(rect)
        
        # Prepare label text
        label_parts = []
        if self.show_labels:
            label_parts.append(class_name)
        if self.show_confidence:
            label_parts.append(f"{confidence:.2f}")
        if self.show_tracking_id and track_id is not None:
            label_parts.append(f"ID:{track_id}")
        
        if not label_parts:
            return
        
        label_text = " | ".join(label_parts)
        
        # Calculate label background
        metrics = painter.fontMetrics()
        label_width = metrics.horizontalAdvance(label_text) + 10
        label_height = metrics.height() + 6
        
        # Draw label background
        label_rect = QRect(int(x1), int(y1-label_height), label_width, label_height)
        painter.fillRect(label_rect, color)
        
        # Draw label text
        painter.setPen(QPen(QColor(255, 255, 255), 1))
        text_rect = QRect(int(x1+5), int(y1-label_height+3), label_width-10, label_height-6)
        painter.drawText(text_rect, Qt.AlignLeft | Qt.AlignVCenter, label_text)
    
    def mousePressEvent(self, event):
        """Handle mouse clicks on detections"""
        if event.button() == Qt.LeftButton:
            click_pos = event.pos()
            
            # Check if click is within any detection bbox
            for detection in self.detections:
                bbox = detection.get('bbox', [0, 0, 100, 100])
                x1, y1, x2, y2 = bbox
                
                if x1 <= click_pos.x() <= x2 and y1 <= click_pos.y() <= y2:
                    self.detection_clicked.emit(detection)
                    break
    
    def clear_detections(self):
        """Clear all detections"""
        self.detections = []
        self.update()
    
    def get_detection_count(self):
        """Get number of current detections"""
        return len(self.detections)
    
    def get_class_counts(self):
        """Get count of detections by class"""
        counts = {}
        for detection in self.detections:
            class_name = detection.get('class', 'unknown')
            counts[class_name] = counts.get(class_name, 0) + 1
        return counts