Update app.py

app.py

@@ -180,24 +180,30 @@ 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 with batch processing for faster performance.
+    Processes the IP camera stream to count objects of the selected classes crossing the line.
     """
     global line_params
 
-    if line_params is None or selected_classes is None or not stream_url:
-        return None, "Error: Missing required parameters"
+    errors = []
+
+    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():
-        return None, "Error: Could not open stream"
-
-    # Initialize variables
-    frames_buffer = []
-    crossed_objects = {}
-    batch_size = 16
-    max_tracked_objects = 1000
+        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)
@@ -205,58 +211,112 @@ def process_video(confidence_threshold=0.5, selected_classes=None, stream_url=No
     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
 
+    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
 
         frames_buffer.append(frame)
 
         if len(frames_buffer) >= batch_size:
             # Process batch of frames
-            results = model.track(frames_buffer, persist=True, conf=confidence_threshold, verbose=False)
+            results = model.track(frames_buffer, persist=True, conf=confidence_threshold)
 
-            # Process each frame's results
-            for frame_idx, result in enumerate(results):
+            # 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()
 
-                    # Create annotated frame
-                    annotated_frame = frames_buffer[frame_idx].copy()
-
                     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
 
-                                # Clear if too many objects
                                 if len(crossed_objects) > max_tracked_objects:
                                     crossed_objects.clear()
 
-                            # Draw bounding box
-                            x1, y1, x2, y2 = map(int, box)
-                            cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+                # Visualize the results with bounding boxes, masks, and IDs
+                annotated_frame = result.plot()
 
-                    # Draw line
-                    draw_angled_line(annotated_frame, line_params, color=(0, 255, 0), thickness=2)
+                # Draw the angled line on the frame
+                draw_angled_line(annotated_frame, line_params, color=(0, 255, 0), thickness=2)
 
-                    # Draw count
-                    count = len(crossed_objects)
-                    cv2.putText(annotated_frame, f"COUNT: {count}", (10, 30), 
-                              cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 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)
 
-                    # Yield the processed frame
-                    yield annotated_frame, ""
+                # 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, ""
 
-            # 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>")