first init

app.py

@@ -0,0 +1,68 @@
+import streamlit as st
+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 streamlit_webrtc import webrtc_streamer, RTCConfiguration
+from video_processor import VideoProcessor
+from utils import process_image, process_video_realtime
+import tempfile
+import numpy as np
+import cv2
+
+RTC_CONFIGURATION = RTCConfiguration({"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]})
+
+def main():
+    st.title('Vehicle Detection and Counting')
+
+    model = YOLO("yolov8n.pt")
+    model.fuse()
+    label_map = model.model.names
+    byte_tracker = ByteTrack()
+    box_annotator = BoxAnnotator(color=ColorPalette.default(), thickness=4, text_thickness=4, text_scale=2)
+
+    source_option = st.selectbox("Select Source Data", ("Webcam", "Image file", "Video file"))
+
+    if source_option == "Webcam":
+        ctx = webrtc_streamer(
+            key="example",
+            video_processor_factory=lambda: VideoProcessor(model, label_map, byte_tracker, box_annotator),
+            rtc_configuration=RTC_CONFIGURATION,
+            media_stream_constraints={"video": True, "audio": False}
+        )
+        if st.button('Take Screenshot and Process'):
+            if ctx.video_processor and ctx.video_processor.frame is not None:
+                image = ctx.video_processor.frame
+                processed_image, object_counts = process_image(image, model, label_map, byte_tracker, box_annotator)
+                st.image(processed_image, caption="Processed Image", use_column_width=True)
+                st.write("Results:", object_counts)
+            else:
+                st.warning("No frame available yet.")
+    elif source_option == "Image file":
+        uploaded_image = st.file_uploader("Upload an image file", type=["jpg", "jpeg", "png"])
+        if uploaded_image is not None:
+            image = cv2.imdecode(np.frombuffer(uploaded_image.read(), np.uint8), cv2.IMREAD_COLOR)
+            st.image(image, caption="Uploaded Image", use_column_width=True)
+            if st.button('Process Image'):
+                processed_image, object_counts = process_image(image, model, label_map, byte_tracker, box_annotator)
+                st.image(processed_image, caption="Processed Image", use_column_width=True)
+                st.write("Results:", object_counts)
+    elif source_option == "Video file":
+        uploaded_video = st.file_uploader("Upload a video file", type=["mp4"])
+        demo_video = "vehicle-counting.mp4"
+        
+        if uploaded_video is not None:
+            temp_video_path = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
+            with open(temp_video_path, 'wb') as f:
+                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)
+        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)
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,5 @@
+ultralytics==8.0.157
+supervision==0.13.0
+tqdm
+streamlit==1.36.0
+streamlit_webrtc==0.47.7

utils.py

@@ -0,0 +1,34 @@
+from supervision.detection.core import Detections
+from supervision.utils.video import VideoInfo, get_video_frames_generator
+import streamlit as st
+
+def process_image(image, model, label_map, byte_tracker, box_annotator):
+    results = model(image)[0]
+    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]
+    annotated_image = box_annotator.annotate(scene=image, detections=detections, labels=labels)
+    
+    object_counts = {}
+    for detection in detections:
+        class_id = detection[3]
+        class_name = label_map[class_id]
+        if class_name in object_counts:
+            object_counts[class_name] += 1
+        else:
+            object_counts[class_name] = 1
+
+    return annotated_image, object_counts
+
+def process_video_realtime(input_video_path, model, label_map, byte_tracker, box_annotator):
+    video_info = VideoInfo.from_video_path(input_video_path)
+    generator = get_video_frames_generator(input_video_path)
+    stframe = st.empty()
+
+    for frame in generator:
+        results = model(frame)[0]
+        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]
+        annotated_frame = box_annotator.annotate(scene=frame, detections=detections, labels=labels)
+        stframe.image(annotated_frame, channels="BGR")

video_processor.py

@@ -0,0 +1,13 @@
+from streamlit_webrtc import VideoProcessorBase
+
+class VideoProcessor(VideoProcessorBase):
+    def __init__(self, model, label_map, byte_tracker, box_annotator):
+        self.model = model
+        self.label_map = label_map
+        self.byte_tracker = byte_tracker
+        self.box_annotator = box_annotator
+        self.frame = None
+
+    def recv(self, frame):
+        self.frame = frame.to_ndarray(format="bgr24")
+        return self.frame