Update app.py

app.py

@@ -3,72 +3,33 @@ import streamlit as st
 import numpy as np
 import tempfile
 import os
+import asyncio
 from ultralytics import YOLO
-from streamlit_webrtc import  (webrtc_streamer, VideoProcessorBase, WebRtcMode, RTCConfiguration)
-import av
-from turn import get_ice_servers
-
-model = YOLO('yolov8n.pt') 
-
-# Global variable to store the latest frame with bounding boxes
-cached_frame = None
-frame_skip = 5  # Process every 5th frame
-# # Define a custom video processor class inheriting from VideoProcessorBase
-# class VideoProcessor(VideoProcessorBase):
-#     def __init__(self):
-#         self.model = model
-#         self.frame_skip = 10  # Class-level variable for frame skipping
-#         self.cached_frame = None  # Class-level variable for cached frames
-        
-def recv(frame: av.VideoFrame) -> av.VideoFrame:
-                # Skip frames to reduce processing load
-                # global frame_skip, cached_frame
-                
-                # if frame_skip > 0:
-                #     frame_skip -= 1
-                #     return frame
-        
-                # Reset frame skip
-                # frame_skip = 5
-        
-                # Convert frame to OpenCV format (BGR)
-                frame_bgr = frame.to_ndarray(format="bgr24")
-        
-                # Resize frame to reduce processing time
-                frame_resized = cv2.resize(frame_bgr, (160, 120)) # Instead of 640x480
-        
-                # # Detect and track objects using YOLOv8
-                # results = model.track(frame_resized, persist=True)
-        
-                # # Plot results
-                # frame_annotated = results[0].plot()
-        
-                # # Cache the annotated frame
-                # cached_frame = frame_annotated
-        
-        
-                # Process every nth frame
-                if  frame_skip == 0:
-                    # Reset the frame skip counter
-                    frame_skip = 10
-        
-                    # Detect and track objects using YOLOv8
-                    results = model.track(frame_resized, persist=True)
-        
-                    # Plot results
-                    frame_annotated = results[0].plot()
-        
-                    # Cache the annotated frame
-                    cached_frame = frame_annotated
-                else:
-                    # Use the cached frame for skipped frames
-                    frame_annotated = cached_frame if cached_frame is not None else frame_resized
-                    frame_skip -= 1
-        
-                # Convert frame back to RGB format
-                frame_rgb = cv2.cvtColor(frame_annotated, cv2.COLOR_BGR2RGB)
-        
-                return av.VideoFrame.from_ndarray(frame_rgb, format="rgb24")
+from streamlit_webrtc import VideoTransformerBase, webrtc_streamer
+
+# Define a video transformer for object tracking
+class ObjectTrackingTransformer(VideoTransformerBase):
+    def __init__(self):
+        # Load YOLOv8 model
+        self.model = YOLO('yolov8n.pt')
+
+    def transform(self, frame):
+        # Convert frame to OpenCV format (BGR)
+        frame_bgr = np.array(frame.to_image())
+
+        # Resize frame to reduce processing time
+        frame_resized = cv2.resize(frame_bgr, (640, 480))
+
+        # Detect and track objects using YOLOv8
+        results = self.model.track(frame_resized, persist=True)
+
+        # Plot results
+        frame_annotated = results[0].plot()
+
+        # Convert frame back to RGB format
+        frame_rgb = cv2.cvtColor(frame_annotated, cv2.COLOR_BGR2RGB)
+
+        return frame_rgb
 
 # Streamlit web app
 def main():
@@ -83,22 +44,7 @@ def main():
 
     if option == "Live Stream":
         # Start the WebRTC stream with object tracking
-        # WebRTC streamer configuration
-        # Define RTC configuration for WebRTC
-        # RTC_CONFIGURATION = RTCConfiguration({
-        #     "iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]
-        # })
-        # Start the WebRTC stream with object tracking
-        # webrtc_streamer(key="live-stream", video_frame_callback=recv,
-        #                 rtc_configuration=rtc_configuration, sendback_audio=False)
-        webrtc_streamer(key="live-stream", 
-                #mode=WebRtcMode.SENDRECV,  
-                video_frame_callback=recv,
-                rtc_configuration={"iceServers": get_ice_servers()},
-                media_stream_constraints={"video": True, "audio": False},
-                async_processing=True)
-
-
+        webrtc_streamer(key="live-stream", video_transformer_factory=ObjectTrackingTransformer)
 
     elif option == "Upload Video":
         # File uploader for video upload
@@ -124,6 +70,8 @@ def main():
 
 # Function to perform object tracking on uploaded video
 def track_uploaded_video(video_file, stop_button, frame_placeholder):
+    # Load YOLOv8 model
+    model = YOLO('yolov8n.pt')
 
     # Create a temporary file to save the uploaded video
     temp_video = tempfile.NamedTemporaryFile(delete=False)