src/annotator.py

"""
Frame Annotation Module
========================

Provides comprehensive video frame annotation capabilities for vehicle detection
visualization including bounding boxes, labels, trajectories, and counting zones.

Authors:
    - Abhay Gupta (0205CC221005)
    - Aditi Lakhera (0205CC221011)
    - Balraj Patel (0205CC221049)
    - Bhumika Patel (0205CC221050)
"""

import numpy as np
import supervision as sv
from typing import List, Optional, Tuple
import logging

logger = logging.getLogger(__name__)


class FrameAnnotator:
    """
    Comprehensive frame annotation system for vehicle detection visualization.
    
    This class manages multiple annotation layers including bounding boxes,
    labels, object trajectories, counting lines, and detection zones.
    """
    
    def __init__(
        self,
        video_resolution: Tuple[int, int],
        show_boxes: bool = True,
        show_labels: bool = True,
        show_traces: bool = True,
        show_line_zones: bool = True,
        trace_length: int = 20,
        zone_polygon: Optional[np.ndarray] = None
    ):
        """
        Initialize the frame annotator.
        
        Args:
            video_resolution: Video resolution as (width, height)
            show_boxes: Enable bounding box annotations
            show_labels: Enable label annotations
            show_traces: Enable trajectory trace annotations
            show_line_zones: Enable line zone annotations
            trace_length: Maximum number of points in trajectory trace
            zone_polygon: Optional polygon defining detection zone
        """
        self.resolution = video_resolution
        self.show_boxes = show_boxes
        self.show_labels = show_labels
        self.show_traces = show_traces
        self.show_line_zones = show_line_zones
        self.zone_polygon = zone_polygon
        
        # Calculate optimal annotation parameters based on resolution
        self.line_thickness = sv.calculate_optimal_line_thickness(video_resolution)
        self.text_scale = sv.calculate_optimal_text_scale(video_resolution)
        
        logger.info(f"Initializing annotator for {video_resolution[0]}x{video_resolution[1]}")
        logger.info(f"Line thickness: {self.line_thickness}, Text scale: {self.text_scale}")
        
        # Initialize annotation components
        self._setup_annotators(trace_length)
    
    def _setup_annotators(self, trace_length: int) -> None:
        """
        Set up individual annotation components.
        
        Args:
            trace_length: Maximum trajectory trace length
        """
        # Bounding box annotator
        if self.show_boxes:
            self.box_drawer = sv.BoxAnnotator(
                thickness=self.line_thickness,
                color_lookup=sv.ColorLookup.TRACK
            )
            logger.debug("Box annotator initialized")
        else:
            self.box_drawer = None
        
        # Label annotator
        if self.show_labels:
            self.label_drawer = sv.LabelAnnotator(
                text_thickness=self.line_thickness,
                text_scale=self.text_scale,
                text_position=sv.Position.BOTTOM_CENTER,
                color_lookup=sv.ColorLookup.TRACK
            )
            logger.debug("Label annotator initialized")
        else:
            self.label_drawer = None
        
        # Trajectory trace annotator
        if self.show_traces:
            self.trace_drawer = sv.TraceAnnotator(
                thickness=self.line_thickness,
                trace_length=trace_length,
                position=sv.Position.BOTTOM_CENTER,
                color_lookup=sv.ColorLookup.TRACK
            )
            logger.debug(f"Trace annotator initialized with length {trace_length}")
        else:
            self.trace_drawer = None
        
        # Line zone annotator
        if self.show_line_zones:
            self.line_zone_drawer = sv.LineZoneAnnotator(
                thickness=self.line_thickness,
                text_thickness=self.line_thickness,
                text_scale=self.text_scale,
                color=sv.Color.WHITE,
                text_color=sv.Color.BLACK,
                display_in_count=True,
                display_out_count=True
            )
            logger.debug("Line zone annotator initialized")
        else:
            self.line_zone_drawer = None
        
        # Polygon zone annotator (for detection area visualization)
        if self.zone_polygon is not None:
            try:
                zone = sv.PolygonZone(
                    polygon=self.zone_polygon,
                    frame_resolution_wh=self.resolution
                )
                self.polygon_drawer = sv.PolygonZoneAnnotator(
                    zone=zone,
                    color=sv.Color.GREEN,
                    thickness=self.line_thickness,
                    display_in_zone_count=False
                )
                logger.debug("Polygon zone annotator initialized")
            except Exception as e:
                logger.warning(f"Failed to initialize polygon zone: {e}")
                self.polygon_drawer = None
        else:
            self.polygon_drawer = None
    
    def draw_annotations(
        self,
        frame: np.ndarray,
        detections: sv.Detections,
        labels: Optional[List[str]] = None,
        line_zones: Optional[List[sv.LineZone]] = None
    ) -> np.ndarray:
        """
        Apply all enabled annotations to a video frame.
        
        Args:
            frame: Input video frame (BGR format)
            detections: Detection results to annotate
            labels: Optional list of labels for each detection
            line_zones: Optional list of line zones for counting visualization
            
        Returns:
            Annotated frame
        """
        # Create a copy to avoid modifying original
        annotated = frame.copy()
        
        try:
            # Draw polygon zone first (background layer)
            if self.polygon_drawer is not None:
                annotated = self.polygon_drawer.annotate(annotated)
            
            # Draw bounding boxes
            if self.box_drawer is not None and len(detections) > 0:
                annotated = self.box_drawer.annotate(
                    scene=annotated,
                    detections=detections
                )
            
            # Draw trajectory traces
            if self.trace_drawer is not None and len(detections) > 0:
                annotated = self.trace_drawer.annotate(
                    scene=annotated,
                    detections=detections
                )
            
            # Draw labels
            if self.label_drawer is not None and len(detections) > 0:
                annotated = self.label_drawer.annotate(
                    scene=annotated,
                    detections=detections,
                    labels=labels
                )
            
            # Draw line zones
            if self.line_zone_drawer is not None and line_zones:
                for zone in line_zones:
                    annotated = self.line_zone_drawer.annotate(
                        frame=annotated,
                        line_counter=zone
                    )
            
            return annotated
            
        except Exception as e:
            logger.error(f"Error during annotation: {e}")
            # Return original frame if annotation fails
            return frame
    
    def update_settings(
        self,
        show_boxes: Optional[bool] = None,
        show_labels: Optional[bool] = None,
        show_traces: Optional[bool] = None,
        show_line_zones: Optional[bool] = None
    ) -> None:
        """
        Update annotation settings dynamically.
        
        Args:
            show_boxes: Enable/disable bounding boxes
            show_labels: Enable/disable labels
            show_traces: Enable/disable traces
            show_line_zones: Enable/disable line zones
        """
        if show_boxes is not None:
            self.show_boxes = show_boxes
        if show_labels is not None:
            self.show_labels = show_labels
        if show_traces is not None:
            self.show_traces = show_traces
        if show_line_zones is not None:
            self.show_line_zones = show_line_zones
        
        logger.info(f"Annotation settings updated: boxes={self.show_boxes}, "
                   f"labels={self.show_labels}, traces={self.show_traces}, "
                   f"zones={self.show_line_zones}")