Update app.py

app.py

@@ -14,8 +14,6 @@ cv2.setNumThreads(cpu_cores)
 print(f"OpenCV using {cv2.getNumThreads()} threads out of {cpu_cores} available cores")
 
 ##############
-
-
 import cv2
 import gradio as gr
 import numpy as np
@@ -180,7 +178,6 @@ def draw_angled_line(image, line_params, color=(0, 255, 0), thickness=2):
     _, _, start_point, end_point = line_params
     cv2.line(image, start_point, end_point, color, thickness)
 
-
 def process_video(confidence_threshold=0.5, selected_classes=None, stream_url=None):
     """
     Processes the IP camera stream to count objects of the selected classes crossing the line.
@@ -205,16 +202,9 @@ def process_video(confidence_threshold=0.5, selected_classes=None, stream_url=No
         errors.append("Error: Could not open stream.")
         return None, "\n".join(errors)
 
-    # Set capture properties for better performance
-    cap.set(cv2.CAP_PROP_BUFFERSIZE, 30)
-    cap.set(cv2.CAP_PROP_FPS, 30)
-    cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'MJPG'))
-
     model = YOLO(model="yolo11n.pt")
     crossed_objects = {}
-    max_tracked_objects = 1000
-    frames_buffer = []
-    batch_size = 16
+    max_tracked_objects = 1000  # Maximum number of objects to track before clearing
 
     logger.info("Starting to process the stream...")
     while cap.isOpened():
@@ -223,96 +213,48 @@ def process_video(confidence_threshold=0.5, selected_classes=None, stream_url=No
             errors.append("Error: Could not read frame from the stream.")
             break
 
-        frames_buffer.append(frame)
-
-        if len(frames_buffer) >= batch_size:
-            # Process batch of frames
-            results = model.track(frames_buffer, persist=True, conf=confidence_threshold)
-
-            # Process and yield each frame immediately to maintain real-time appearance
-            for idx, result in enumerate(results):
-                if result.boxes.id is not None:
-                    track_ids = result.boxes.id.int().cpu().tolist()
-                    clss = result.boxes.cls.cpu().tolist()
-                    boxes = result.boxes.xyxy.cpu()
-                    confs = result.boxes.conf.cpu().tolist()
-
-                    for box, cls, t_id, conf in zip(boxes, clss, track_ids, confs):
-                        if conf >= confidence_threshold and model.names[cls] in selected_classes:
-                            if is_object_crossing_line(box, line_params) and t_id not in crossed_objects:
-                                crossed_objects[t_id] = True
-
-                                if len(crossed_objects) > max_tracked_objects:
-                                    crossed_objects.clear()
-
-                # Visualize the results with bounding boxes, masks, and IDs
-                annotated_frame = result.plot()
-
-                # Draw the angled line on the frame
-                draw_angled_line(annotated_frame, line_params, color=(0, 255, 0), thickness=2)
-
-                # Display the count on the frame with a modern look
-                count = len(crossed_objects)
-                (text_width, text_height), _ = cv2.getTextSize(f"COUNT: {count}", cv2.FONT_HERSHEY_SIMPLEX, 1, 2)
-
-                # Calculate the position for the middle of the top
-                margin = 10
-                x = (annotated_frame.shape[1] - text_width) // 2
-                y = text_height + margin
-
-                # Draw the black background rectangle
-                cv2.rectangle(annotated_frame, 
-                            (x - margin, y - text_height - margin), 
-                            (x + text_width + margin, y + margin), 
-                            (0, 0, 0), -1)
-
-                # Draw the text
-                cv2.putText(annotated_frame, f"COUNT: {count}", 
-                           (x, y), cv2.FONT_HERSHEY_SIMPLEX, 
-                           1, (0, 255, 0), 2)
-
-                # Yield each frame as soon as it's processed
-                yield annotated_frame, ""
-
-            # Clear the buffer after processing
-            frames_buffer = []
-        
-        # If we have remaining frames that don't make a full batch, process them too
-        elif frames_buffer:
-            results = model.track(frames_buffer, persist=True, conf=confidence_threshold)
-            
-            for result in results:
-                if result.boxes.id is not None:
-                    track_ids = result.boxes.id.int().cpu().tolist()
-                    clss = result.boxes.cls.cpu().tolist()
-                    boxes = result.boxes.xyxy.cpu()
-                    confs = result.boxes.conf.cpu().tolist()
-
-                    for box, cls, t_id, conf in zip(boxes, clss, track_ids, confs):
-                        if conf >= confidence_threshold and model.names[cls] in selected_classes:
-                            if is_object_crossing_line(box, line_params) and t_id not in crossed_objects:
-                                crossed_objects[t_id] = True
-
-                annotated_frame = result.plot()
-                draw_angled_line(annotated_frame, line_params, color=(0, 255, 0), thickness=2)
-
-                count = len(crossed_objects)
-                (text_width, text_height), _ = cv2.getTextSize(f"COUNT: {count}", cv2.FONT_HERSHEY_SIMPLEX, 1, 2)
-                margin = 10
-                x = (annotated_frame.shape[1] - text_width) // 2
-                y = text_height + margin
-
-                cv2.rectangle(annotated_frame, 
-                            (x - margin, y - text_height - margin), 
-                            (x + text_width + margin, y + margin), 
-                            (0, 0, 0), -1)
-                cv2.putText(annotated_frame, f"COUNT: {count}", 
-                           (x, y), cv2.FONT_HERSHEY_SIMPLEX, 
-                           1, (0, 255, 0), 2)
-
-                yield annotated_frame, ""
-
-            frames_buffer = []
+        # Perform object tracking with confidence threshold
+        results = model.track(frame, persist=True, conf=confidence_threshold)
+
+        if results[0].boxes.id is not None:
+            track_ids = results[0].boxes.id.int().cpu().tolist()
+            clss = results[0].boxes.cls.cpu().tolist()
+            boxes = results[0].boxes.xyxy.cpu()
+            confs = results[0].boxes.conf.cpu().tolist()
+
+            for box, cls, t_id, conf in zip(boxes, clss, track_ids, confs):
+                if conf >= confidence_threshold and model.names[cls] in selected_classes:
+                    # Check if the object crosses the line
+                    if is_object_crossing_line(box, line_params) and t_id not in crossed_objects:
+                        crossed_objects[t_id] = True
+
+                        # Clear the dictionary if it gets too large
+                        if len(crossed_objects) > max_tracked_objects:
+                            crossed_objects.clear()
+
+        # Visualize the results with bounding boxes, masks, and IDs
+        annotated_frame = results[0].plot()
+
+        # Draw the angled line on the frame
+        draw_angled_line(annotated_frame, line_params, color=(0, 255, 0), thickness=2)
+
+        # Display the count on the frame with a modern look
+        count = len(crossed_objects)
+        (text_width, text_height), _ = cv2.getTextSize(f"COUNT: {count}", cv2.FONT_HERSHEY_SIMPLEX, 1, 2)
+
+        # Calculate the position for the middle of the top
+        margin = 10  # Margin from the top
+        x = (annotated_frame.shape[1] - text_width) // 2  # Center-align the text horizontally
+        y = text_height + margin  # Top-align the text
+
+        # Draw the black background rectangle
+        cv2.rectangle(annotated_frame, (x - margin, y - text_height - margin), (x + text_width + margin, y + margin), (0, 0, 0), -1)
+
+        # Draw the text
+        cv2.putText(annotated_frame, f"COUNT: {count}", (x, y), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+
+        # Yield the annotated frame to Gradio
+        yield annotated_frame, ""
 
     cap.release()
     logger.info("Stream processing completed.")