chore: import upstream snapshot with attribution

examples/traffic_analysis/inference_example.py

@@ -0,0 +1,223 @@
+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)