Update src/annotator.py

src/annotator.py

@@ -1,237 +1,234 @@
-"""
-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}")
+"""
+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}")