File size: 8,482 Bytes
fdf832f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 |
"""
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}")
|