""" 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) 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}")