Update app.py

app.py

@@ -1,30 +1,28 @@
+import multiprocessing
 import cv2
 import gradio as gr
 import numpy as np
 from PIL import Image, ImageDraw
 from ultralytics import YOLO
+from ultralytics.utils.plotting import Annotator, colors
 import logging
-import threading
-import queue
+import math
 import time
+from collections import deque
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Global variables for line coordinates and line equation
+# Global variables to store line coordinates and line equation
 start_point = None
 end_point = None
-line_params = None  # Stores (slope, intercept, start_point, end_point)
+line_params = None  # Stores (slope, intercept) of the line
 
-# Low-resolution for inference
-LOW_RES = (320, 180)
-
-# Frame queue for processed frames
-frame_queue = queue.Queue(maxsize=30)  # Adjust queue size based on memory constraints
-
-# Thread control flag
-processing_active = True
+# Maximize CPU usage
+cpu_cores = multiprocessing.cpu_count()
+cv2.setNumThreads(cpu_cores)
+logger.info(f"OpenCV using {cv2.getNumThreads()} threads out of {cpu_cores} available cores")
 
 def extract_first_frame(stream_url):
     """
@@ -56,27 +54,46 @@ def update_line(image, evt: gr.SelectData):
     """
     global start_point, end_point, line_params
 
+    # If it's the first click, set the start point and show it on the image
     if start_point is None:
         start_point = (evt.index[0], evt.index[1])
+
+        # Draw the start point on the image
         draw = ImageDraw.Draw(image)
-        draw.ellipse((start_point[0] - 5, start_point[1] - 5, start_point[0] + 5, start_point[1] + 5), fill="blue", outline="blue")
+        draw.ellipse(
+            (start_point[0] - 5, start_point[1] - 5, start_point[0] + 5, start_point[1] + 5),
+            fill="blue", outline="blue"
+        )
+
         return image, f"Line Coordinates:\nStart: {start_point}, End: None"
 
+    # If it's the second click, set the end point and draw the line
     end_point = (evt.index[0], evt.index[1])
+
+    # Calculate the slope (m) and intercept (b) of the line: y = mx + b
     if start_point[0] != end_point[0]:  # Avoid division by zero
         slope = (end_point[1] - start_point[1]) / (end_point[0] - start_point[0])
         intercept = start_point[1] - slope * start_point[0]
-        line_params = (slope, intercept, start_point, end_point)
+        line_params = (slope, intercept, start_point, end_point)  # Store slope, intercept, and points
     else:
+        # Vertical line (special case)
         line_params = (float('inf'), start_point[0], start_point, end_point)
 
+    # Draw the line and end point on the image
     draw = ImageDraw.Draw(image)
     draw.line([start_point, end_point], fill="red", width=2)
-    draw.ellipse((end_point[0] - 5, end_point[1] - 5, end_point[0] + 5, end_point[1] + 5), fill="green", outline="green")
+    draw.ellipse(
+        (end_point[0] - 5, end_point[1] - 5, end_point[0] + 5, end_point[1] + 5),
+        fill="green", outline="green"
+    )
 
+    # Return the updated image and line info
     line_info = f"Line Coordinates:\nStart: {start_point}, End: {end_point}\nLine Equation: y = {line_params[0]:.2f}x + {line_params[1]:.2f}"
+
+    # Reset the points for the next interaction
     start_point = None
     end_point = None
+
     return image, line_info
 
 def reset_line():
@@ -89,19 +106,6 @@ def reset_line():
     line_params = None
     return None, "Line reset. Click to draw a new line."
 
-def is_object_crossing_line(box, line_params):
-    """
-    Determines if an object's bounding box is fully intersected by the user-drawn line.
-    """
-    _, _, line_start, line_end = line_params
-    x1, y1, x2, y2 = box
-    box_edges = [((x1, y1), (x2, y1)), ((x2, y1), (x2, y2)), ((x2, y2), (x1, y2)), ((x1, y2), (x1, y1))]
-    intersection_count = 0
-    for edge_start, edge_end in box_edges:
-        if intersect(line_start, line_end, edge_start, edge_end):
-            intersection_count += 1
-    return intersection_count >= 2
-
 def intersect(A, B, C, D):
     """
     Determines if two line segments AB and CD intersect.
@@ -110,51 +114,54 @@ def intersect(A, B, C, D):
         return (C[1] - A[1]) * (B[0] - A[0]) - (B[1] - A[1]) * (C[0] - A[0])
 
     def on_segment(A, B, C):
-        return min(A[0], B[0]) <= C[0] <= max(A[0], B[0]) and min(A[1], B[1]) <= C[1] <= max(A[1], B[1])
+        if min(A[0], B[0]) <= C[0] <= max(A[0], B[0]) and min(A[1], B[1]) <= C[1] <= max(A[1], B[1]):
+            return True
+        return False
 
+    # Check if the line segments intersect
     ccw1 = ccw(A, B, C)
     ccw2 = ccw(A, B, D)
     ccw3 = ccw(C, D, A)
     ccw4 = ccw(C, D, B)
-    return ((ccw1 * ccw2 < 0) and (ccw3 * ccw4 < 0)) or (ccw1 == 0 and on_segment(A, B, C)) or (ccw2 == 0 and on_segment(A, B, D)) or (ccw3 == 0 and on_segment(C, D, A)) or (ccw4 == 0 and on_segment(C, D, B))
 
-def process_frames(stream_url, confidence_threshold, selected_classes):
+    if ((ccw1 * ccw2 < 0) and (ccw3 * ccw4 < 0)):
+        return True
+    elif ccw1 == 0 and on_segment(A, B, C):
+        return True
+    elif ccw2 == 0 and on_segment(A, B, D):
+        return True
+    elif ccw3 == 0 and on_segment(C, D, A):
+        return True
+    elif ccw4 == 0 and on_segment(C, D, B):
+        return True
+    else:
+        return False
+
+def is_object_crossing_line(box, line_params):
     """
-    Processes frames in a separate thread and adds them to the frame queue.
+    Determines if an object's bounding box is fully intersected by the user-drawn line.
     """
-    global processing_active, frame_queue
-    cap = cv2.VideoCapture(stream_url)
-    model = YOLO(model="yolo11n.pt")
-    crossed_objects = {}
+    _, _, line_start, line_end = line_params
 
-    while processing_active and cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
+    # Get the bounding box coordinates
+    x1, y1, x2, y2 = box
 
-        # Perform detection on low-res frame
-        low_res_frame = cv2.resize(frame, LOW_RES)
-        results = model.track(low_res_frame, persist=True, conf=confidence_threshold)
-
-        # Scale bounding boxes to high-res
-        scale_x = frame.shape[1] / LOW_RES[0]
-        scale_y = frame.shape[0] / LOW_RES[1]
-        for detection in results[0].boxes.data:
-            x1, y1, x2, y2, conf, cls = detection
-            x1, y1, x2, y2 = int(x1 * scale_x), int(y1 * scale_y), int(x2 * scale_x), int(y2 * scale_y)
-            if is_object_crossing_line((x1, y1, x2, y2), line_params):
-                crossed_objects[results[0].boxes.id.int().cpu().tolist()[0]] = True
-
-        # Draw bounding boxes and line on the frame
-        annotated_frame = results[0].plot()
-        if line_params:
-            draw_angled_line(annotated_frame, line_params, color=(0, 255, 0), thickness=2)
+    # Define the four edges of the bounding box
+    box_edges = [
+        ((x1, y1), (x2, y1)),  # Top edge
+        ((x2, y1), (x2, y2)),  # Right edge
+        ((x2, y2), (x1, y2)),  # Bottom edge
+        ((x1, y2), (x1, y1))   # Left edge
+    ]
 
-        # Add frame to the queue
-        if not frame_queue.full():
-            frame_queue.put(annotated_frame)
+    # Count the number of intersections between the line and the bounding box edges
+    intersection_count = 0
+    for edge_start, edge_end in box_edges:
+        if intersect(line_start, line_end, edge_start, edge_end):
+            intersection_count += 1
 
-    cap.release()
+    # Only count the object if the line intersects the bounding box at least twice
+    return intersection_count >= 2
 
 def draw_angled_line(image, line_params, color=(0, 255, 0), thickness=2):
     """
@@ -163,22 +170,102 @@ 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 display_frames():
+def process_video(confidence_threshold=0.5, selected_classes=None, stream_url=None):
     """
-    Displays frames from the queue at a consistent frame rate.
+    Processes the IP camera stream to count objects of the selected classes crossing the line.
     """
-    while processing_active:
-        if not frame_queue.empty():
-            frame = frame_queue.get()
-            yield cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), ""
-        else:
-            time.sleep(0.03)  # Wait for the next frame
+    global line_params
+
+    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():
+        errors.append("Error: Could not open stream.")
+        return None, "\n".join(errors)
+
+    model = YOLO(model="yolov8n.pt")
+    crossed_objects = {}
+    max_tracked_objects = 1000  # Maximum number of objects to track before clearing
+
+    # Queue to hold frames for processing
+    frame_queue = deque(maxlen=10)
+
+    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
+
+        # Add frame to the queue
+        frame_queue.append(frame)
+
+        # Process frames in the queue
+        if len(frame_queue) > 0:
+            process_frame = frame_queue.popleft()
+
+            # Perform object tracking with confidence threshold
+            results = model.track(process_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.")
 
 # 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>")
     gr.Markdown("## https://github.com/SanshruthR/CCTV_SENTRY_YOLO11")
-
+    
     # Step 1: Enter the IP Camera Stream URL
     stream_url = gr.Textbox(label="Enter IP Camera Stream URL", value="https://s104.ipcamlive.com/streams/68idokwtondsqpmkr/stream.m3u8", visible=False)
 
@@ -188,36 +275,33 @@ with gr.Blocks() as demo:
     if first_frame is None:
         gr.Markdown(f"**Error:** {status}")
     else:
+        # Image component for displaying the first frame
         image = gr.Image(value=first_frame, label="First Frame of Stream", type="pil")
+
         line_info = gr.Textbox(label="Line Coordinates", value="Line Coordinates:\nStart: None, End: None")
         image.select(update_line, inputs=image, outputs=[image, line_info])
 
         # Step 2: Select classes to detect
         gr.Markdown("### Step 2: Select Classes to Detect")
-        model = YOLO(model="yolo11n.pt")
-        class_names = list(model.names.values())
+        model = YOLO(model="yolov8n.pt")  # Load the model to get class names
+        class_names = list(model.names.values())  # Get class names
         selected_classes = gr.CheckboxGroup(choices=class_names, label="Select Classes to Detect")
 
-        # Step 3: Adjust confidence threshold
+        # Step 3: Adjust confidence threshold 
         gr.Markdown("### Step 3: Adjust Confidence Threshold (Optional)")
         confidence_threshold = gr.Slider(minimum=0.0, maximum=1.0, value=0.2, label="Confidence Threshold")
 
         # Process the stream
         process_button = gr.Button("Process Stream")
+
+        # Output image for real-time frame rendering
         output_image = gr.Image(label="Processed Frame", streaming=True)
+
+        # Error box to display warnings/errors
         error_box = gr.Textbox(label="Errors/Warnings", interactive=False)
 
         # Event listener for processing the video
-        process_button.click(
-            fn=lambda: (setattr(globals(), "processing_active", True), threading.Thread(target=process_frames, args=(stream_url.value, confidence_threshold.value, selected_classes.value)).start()),
-            outputs=None
-        )
-
-        # Display frames using a custom thread
-        def start_display_thread():
-            threading.Thread(target=display_frames, daemon=True).start()
-
-        demo.load(start_display_thread, inputs=None, outputs=[output_image, error_box])
+        process_button.click(process_video, inputs=[confidence_threshold, selected_classes, stream_url], outputs=[output_image, error_box])
 
 # Launch the interface
 demo.launch(debug=True)