import os from collections.abc import Iterable import cv2 import numpy as np from inference.models.utils import get_roboflow_model from tqdm import tqdm import supervision as sv COLORS = sv.ColorPalette.from_hex(["#E6194B", "#3CB44B", "#FFE119", "#3C76D1"]) ZONE_IN_POLYGONS = [ np.array([[592, 282], [900, 282], [900, 82], [592, 82]]), np.array([[950, 860], [1250, 860], [1250, 1060], [950, 1060]]), np.array([[592, 582], [592, 860], [392, 860], [392, 582]]), np.array([[1250, 282], [1250, 530], [1450, 530], [1450, 282]]), ] ZONE_OUT_POLYGONS = [ np.array([[950, 282], [1250, 282], [1250, 82], [950, 82]]), np.array([[592, 860], [900, 860], [900, 1060], [592, 1060]]), np.array([[592, 282], [592, 550], [392, 550], [392, 282]]), np.array([[1250, 860], [1250, 560], [1450, 560], [1450, 860]]), ] class DetectionsManager: def __init__(self) -> None: self.tracker_id_to_zone_id: dict[int, int] = {} self.counts: dict[int, dict[int, set[int]]] = {} def update( self, detections_all: sv.Detections, detections_in_zones: list[sv.Detections], detections_out_zones: list[sv.Detections], ) -> sv.Detections: for zone_in_id, detections_in_zone in enumerate(detections_in_zones): for tracker_id in detections_in_zone.tracker_id: self.tracker_id_to_zone_id.setdefault(tracker_id, zone_in_id) for zone_out_id, detections_out_zone in enumerate(detections_out_zones): for tracker_id in detections_out_zone.tracker_id: if tracker_id in self.tracker_id_to_zone_id: zone_in_id = self.tracker_id_to_zone_id[tracker_id] self.counts.setdefault(zone_out_id, {}) self.counts[zone_out_id].setdefault(zone_in_id, set()) self.counts[zone_out_id][zone_in_id].add(tracker_id) if len(detections_all) > 0: detections_all.class_id = np.vectorize( lambda x: self.tracker_id_to_zone_id.get(x, -1) )(detections_all.tracker_id) else: detections_all.class_id = np.array([], dtype=int) return detections_all[detections_all.class_id != -1] def initiate_polygon_zones( polygons: list[np.ndarray], triggering_anchors: Iterable[sv.Position] = [sv.Position.CENTER], ) -> list[sv.PolygonZone]: return [ sv.PolygonZone( polygon=polygon, triggering_anchors=triggering_anchors, ) for polygon in polygons ] class VideoProcessor: def __init__( self, roboflow_api_key: str, model_id: str, source_video_path: str, target_video_path: str | None = None, confidence_threshold: float = 0.3, iou_threshold: float = 0.7, ) -> None: self.conf_threshold = confidence_threshold self.iou_threshold = iou_threshold self.source_video_path = source_video_path self.target_video_path = target_video_path self.model = get_roboflow_model(model_id=model_id, api_key=roboflow_api_key) self.tracker = sv.ByteTrack() self.video_info = sv.VideoInfo.from_video_path(source_video_path) self.zones_in = initiate_polygon_zones(ZONE_IN_POLYGONS, [sv.Position.CENTER]) self.zones_out = initiate_polygon_zones(ZONE_OUT_POLYGONS, [sv.Position.CENTER]) self.box_annotator = sv.BoxAnnotator(color=COLORS) self.label_annotator = sv.LabelAnnotator( color=COLORS, text_color=sv.Color.BLACK ) self.trace_annotator = sv.TraceAnnotator( color=COLORS, position=sv.Position.CENTER, trace_length=100, thickness=2 ) self.detections_manager = DetectionsManager() def process_video(self) -> None: frame_generator = sv.get_video_frames_generator( source_path=self.source_video_path ) if self.target_video_path: with sv.VideoSink(self.target_video_path, self.video_info) as sink: for frame in tqdm(frame_generator, total=self.video_info.total_frames): annotated_frame = self.process_frame(frame) sink.write_frame(annotated_frame) else: for frame in tqdm(frame_generator, total=self.video_info.total_frames): annotated_frame = self.process_frame(frame) cv2.imshow("Processed Video", annotated_frame) if cv2.waitKey(1) & 0xFF == ord("q"): break cv2.destroyAllWindows() def annotate_frame( self, frame: np.ndarray, detections: sv.Detections ) -> np.ndarray: annotated_frame = frame.copy() for i, (zone_in, zone_out) in enumerate(zip(self.zones_in, self.zones_out)): annotated_frame = sv.draw_polygon( annotated_frame, zone_in.polygon, COLORS.colors[i] ) annotated_frame = sv.draw_polygon( annotated_frame, zone_out.polygon, COLORS.colors[i] ) labels = [f"#{tracker_id}" for tracker_id in detections.tracker_id] annotated_frame = self.trace_annotator.annotate(annotated_frame, detections) annotated_frame = self.box_annotator.annotate(annotated_frame, detections) annotated_frame = self.label_annotator.annotate( annotated_frame, detections, labels ) for zone_out_id, zone_out in enumerate(self.zones_out): zone_center = sv.get_polygon_center(polygon=zone_out.polygon) if zone_out_id in self.detections_manager.counts: counts = self.detections_manager.counts[zone_out_id] for i, zone_in_id in enumerate(counts): count = len(self.detections_manager.counts[zone_out_id][zone_in_id]) text_anchor = sv.Point(x=zone_center.x, y=zone_center.y + 40 * i) annotated_frame = sv.draw_text( scene=annotated_frame, text=str(count), text_anchor=text_anchor, background_color=COLORS.colors[zone_in_id], ) return annotated_frame def process_frame(self, frame: np.ndarray) -> np.ndarray: results = self.model.infer( frame, confidence=self.conf_threshold, iou_threshold=self.iou_threshold )[0] detections = sv.Detections.from_inference(results) detections.class_id = np.zeros(len(detections)) detections = self.tracker.update_with_detections(detections) detections_in_zones = [] detections_out_zones = [] for zone_in, zone_out in zip(self.zones_in, self.zones_out): detections_in_zone = detections[zone_in.trigger(detections=detections)] detections_in_zones.append(detections_in_zone) detections_out_zone = detections[zone_out.trigger(detections=detections)] detections_out_zones.append(detections_out_zone) detections = self.detections_manager.update( detections, detections_in_zones, detections_out_zones ) return self.annotate_frame(frame, detections) def main( source_video_path: str, target_video_path: str, roboflow_api_key: str, model_id: str = "vehicle-count-in-drone-video/6", confidence_threshold: float = 0.3, iou_threshold: float = 0.7, ) -> None: """ Traffic Flow Analysis with Inference and ByteTrack. Args: source_video_path: Path to the source video file target_video_path: Path to the target video file (output) roboflow_api_key: Roboflow API key model_id: Roboflow model ID confidence_threshold: Confidence threshold for the model iou_threshold: IOU threshold for the model """ api_key = roboflow_api_key api_key = os.environ.get("ROBOFLOW_API_KEY", api_key) if api_key is None: raise ValueError( "Roboflow API KEY is missing. Please provide it as an argument or set the " "ROBOFLOW_API_KEY environment variable." ) roboflow_api_key = api_key processor = VideoProcessor( roboflow_api_key=roboflow_api_key, model_id=model_id, source_video_path=source_video_path, target_video_path=target_video_path, confidence_threshold=confidence_threshold, iou_threshold=iou_threshold, ) processor.process_video() if __name__ == "__main__": from jsonargparse import auto_cli, set_parsing_settings set_parsing_settings(parse_optionals_as_positionals=True) auto_cli(main, as_positional=False)