Update app.py

app.py

@@ -182,85 +182,81 @@ def draw_angled_line(image, line_params, color=(0, 255, 0), thickness=2):
 
 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.
+    Processes the IP camera stream with batch processing for faster performance.
     """
     global line_params
 
-    errors = []
+    if line_params is None or selected_classes is None or not stream_url:
+        return None, "Error: Missing required parameters"
 
-    if line_params is None:
-        errors.append("Error: No line drawn. Please draw a line on the first frame.")
-    if selected_classes is None or len(selected_classes) == 0:
-        errors.append("Error: No classes selected. Please select at least one class to detect.")
-    if stream_url is None or stream_url.strip() == "":
-        errors.append("Error: No stream URL provided.")
-
-    if errors:
-        return None, "\n".join(errors)
-
-    logger.info("Connecting to the IP camera stream...")
     cap = cv2.VideoCapture(stream_url)
     if not cap.isOpened():
-        errors.append("Error: Could not open stream.")
-        return None, "\n".join(errors)
+        return None, "Error: Could not open stream"
 
-    model = YOLO(model="yolo11n.pt")
+    # Initialize variables
+    frames_buffer = []
     crossed_objects = {}
-    max_tracked_objects = 100  # Maximum number of objects to track before clearing
+    batch_size = 16
+    max_tracked_objects = 1000
+
+    # 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")
 
-    logger.info("Starting to process the stream...")
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
-            errors.append("Error: Could not read frame from the stream.")
             break
 
-        # Perform object tracking with confidence threshold
-        results = model.track(frame, persist=True, conf=confidence_threshold,iou=0.5,max_det=50,verbose=False)
+        frames_buffer.append(frame)
 
-        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()
+        if len(frames_buffer) >= batch_size:
+            # Process batch of frames
+            results = model.track(frames_buffer, persist=True, conf=confidence_threshold, verbose=False)
 
-            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
+            # Process each frame's results
+            for frame_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()
 
-                        # Clear the dictionary if it gets too large
-                        if len(crossed_objects) > max_tracked_objects:
-                            crossed_objects.clear()
+                    # Create annotated frame
+                    annotated_frame = frames_buffer[frame_idx].copy()
 
-        # Visualize the results with bounding boxes, masks, and IDs
-        annotated_frame = results[0].plot()
+                    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 line crossing
+                            if is_object_crossing_line(box, line_params) and t_id not in crossed_objects:
+                                crossed_objects[t_id] = True
 
-        # Draw the angled line on the frame
-        draw_angled_line(annotated_frame, line_params, color=(0, 255, 0), thickness=2)
+                                # Clear if too many objects
+                                if len(crossed_objects) > max_tracked_objects:
+                                    crossed_objects.clear()
 
-        # 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)
+                            # Draw bounding box
+                            x1, y1, x2, y2 = map(int, box)
+                            cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (0, 255, 0), 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 line
+                    draw_angled_line(annotated_frame, line_params, color=(0, 255, 0), thickness=2)
 
-        # 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 count
+                    count = len(crossed_objects)
+                    cv2.putText(annotated_frame, f"COUNT: {count}", (10, 30), 
+                              cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
 
-        # Draw the text
-        cv2.putText(annotated_frame, f"COUNT: {count}", (x, y), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+                    # Yield the processed frame
+                    yield annotated_frame, ""
 
-        # Yield the annotated frame to Gradio
-        yield annotated_frame, ""
+            # Clear buffer after processing batch
+            frames_buffer = []
 
     cap.release()
-    logger.info("Stream processing completed.")
-
 # Define the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("<h1>Real-time monitoring, object tracking, and line-crossing detection for CCTV camera streams.</h1></center>")