feat: add line counter

app.py

@@ -3,6 +3,8 @@ from ultralytics import YOLO
 from supervision.draw.color import ColorPalette
 from supervision.tracker.byte_tracker.core import ByteTrack
 from supervision.detection.annotate import BoxAnnotator
+from supervision.detection.line_counter import LineZone, LineZoneAnnotator
+from supervision.geometry.core import Point
 from streamlit_webrtc import webrtc_streamer, RTCConfiguration
 from video_processor import VideoProcessor
 from utils import process_image, process_video_realtime
@@ -20,6 +22,11 @@ def main():
     label_map = model.model.names
     byte_tracker = ByteTrack()
     box_annotator = BoxAnnotator(color=ColorPalette.default(), thickness=4, text_thickness=4, text_scale=2)
+    
+    line_start = Point(50, 1500)
+    line_end = Point(3840 - 50, 1500)
+    line_counter = LineZone(start=line_start, end=line_end)
+    line_annotator = LineZoneAnnotator(thickness=4, text_thickness=4, text_scale=2)
 
     source_option = st.selectbox("Select Source Data", ("Webcam", "Image file", "Video file"))
 
@@ -57,12 +64,12 @@ def main():
                 f.write(uploaded_video.read())
             st.video(temp_video_path)
             if st.button('Process Uploaded Video'):
-                process_video_realtime(temp_video_path, model, label_map, byte_tracker, box_annotator)
+                process_video_realtime(temp_video_path, model, label_map, byte_tracker, box_annotator, line_counter, line_annotator)
         else:
             st.header("Demo Video Preview")
             st.video(demo_video)
             if st.button('Process Demo Video'):
-                process_video_realtime(demo_video, model, label_map, byte_tracker, box_annotator)
+                process_video_realtime(demo_video, model, label_map, byte_tracker, box_annotator, line_counter, line_annotator)
 
 if __name__ == "__main__":
-    main()
+    main()

utils.py

@@ -20,7 +20,7 @@ def process_image(image, model, label_map, byte_tracker, box_annotator):
 
     return annotated_image, object_counts
 
-def process_video_realtime(input_video_path, model, label_map, byte_tracker, box_annotator):
+def process_video_realtime(input_video_path, model, label_map, byte_tracker, box_annotator, line_counter, line_annotator):
     video_info = VideoInfo.from_video_path(input_video_path)
     generator = get_video_frames_generator(input_video_path)
     stframe = st.empty()
@@ -30,5 +30,7 @@ def process_video_realtime(input_video_path, model, label_map, byte_tracker, box
         detections = Detections.from_ultralytics(results)
         detections = byte_tracker.update_with_detections(detections=detections)
         labels = [f"{label_map[class_id]} {confidence:0.2f} -track_id:{tracker_id}" for _, _, confidence, class_id, tracker_id in detections]
+        line_counter.trigger(detections=detections)
         annotated_frame = box_annotator.annotate(scene=frame, detections=detections, labels=labels)
-        stframe.image(annotated_frame, channels="BGR")
+        line_annotator.annotate(frame=annotated_frame, line_counter=line_counter)
+        stframe.image(annotated_frame, channels="BGR")