Update app.py

app.py

@@ -1,26 +1,10 @@
-# Maximize CPU usage
-import multiprocessing
-import cv2
-
-# Get the number of CPU cores
-cpu_cores = multiprocessing.cpu_count()
-
-# Set OpenCV to use all available cores
-cv2.setNumThreads(cpu_cores)
-
-# Print the number of threads being used (optional)
-print(f"OpenCV using {cv2.getNumThreads()} threads out of {cpu_cores} available cores")
-
-##############
 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 math
-import torch
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -29,199 +13,241 @@ logger = logging.getLogger(__name__)
 # Global variables to store line coordinates and line equation
 start_point = None
 end_point = None
-line_params = None  # Stores (start_point, end_point)
-
-# Load model once globally
-model = YOLO("yolo11n.pt")
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-model = model.to(device)
-
-def liang_barsky(line, bbox):
-    """Optimized line-rectangle intersection check using Liang-Barsky algorithm"""
-    x1, y1 = line[0]
-    x2, y2 = line[1]
-    xmin, ymin, xmax, ymax = bbox
-
-    dx = x2 - x1
-    dy = y2 - y1
-    p = [-dx, dx, -dy, dy]
-    q = [x1 - xmin, xmax - x1, y1 - ymin, ymax - y1]
-    u1 = 0.0
-    u2 = 1.0
-
-    for i in range(4):
-        if p[i] == 0:
-            if q[i] < 0:
-                return False
-            continue
-        t = q[i] / p[i]
-        if p[i] < 0:
-            if t > u1:
-                u1 = t
-        else:
-            if t < u2:
-                u2 = t
-
-    return u1 <= u2
+line_params = None  # Stores (slope, intercept) of the line
 
 def extract_first_frame(stream_url):
-    """Extracts the first available frame from the IP camera stream"""
-    logger.info("Extracting first frame...")
+    """
+    Extracts the first available frame from the IP camera stream and returns it as a PIL image.
+    """
+    logger.info("Attempting to extract the first frame from the stream...")
     cap = cv2.VideoCapture(stream_url)
     if not cap.isOpened():
+        logger.error("Error: Could not open stream.")
         return None, "Error: Could not open stream."
 
     ret, frame = cap.read()
     cap.release()
 
     if not ret:
-        return None, "Error: Could not read frame."
+        logger.error("Error: Could not read the first frame.")
+        return None, "Error: Could not read the first frame."
 
+    # Convert the frame to a PIL image
     frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-    return Image.fromarray(frame_rgb), "First frame extracted successfully."
+    pil_image = Image.fromarray(frame_rgb)
+
+    logger.info("First frame extracted successfully.")
+    return pil_image, "First frame extracted successfully."
 
 def update_line(image, evt: gr.SelectData):
-    """Handles line drawing interactions"""
+    """
+    Updates the line based on user interaction (click and drag).
+    """
     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])
-    line_params = (start_point, end_point)
-    
+
+    # 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)  # 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
-    return image, f"Line Coordinates:\nStart: {line_params[0]}, End: {line_params[1]}"
+    end_point = None
+
+    return image, line_info
 
 def reset_line():
-    """Resets line coordinates"""
+    """
+    Resets the line coordinates.
+    """
     global start_point, end_point, line_params
-    start_point = end_point = line_params = None
+    start_point = None
+    end_point = None
+    line_params = None
     return None, "Line reset. Click to draw a new line."
 
 def is_object_crossing_line(box, line_params):
-    """Optimized line crossing check using Liang-Barsky algorithm"""
-    if not line_params:
-        return False
-    
-    line_start, line_end = line_params
+    """
+    Determines if an object's bounding box is fully intersected by the user-drawn line.
+    """
+    _, _, line_start, line_end = line_params
+
+    # Get the bounding box coordinates
     x1, y1, x2, y2 = box
-    return liang_barsky((line_start, line_end), (x1, y1, x2, y2))
+
+    # 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
+    ]
+
+    # 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
+
+    # 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):
-    """Draws the user-defined line on the frame"""
-    start, end = line_params
-    cv2.line(image, start, end, color, thickness)
+    """
+    Draws the user-defined line on the frame.
+    """
+    _, _, 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):
-    """Main video processing function with optimizations"""
+    """
+    Processes the IP camera stream to count objects of the selected classes crossing the line.
+    """
     global line_params
+
     errors = []
 
-    if not line_params:
-        errors.append("Error: No line drawn.")
-    if not selected_classes:
-        errors.append("Error: No classes selected.")
-    if not stream_url:
+    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)
 
-    # Convert class names to indices once
-    selected_class_indices = {i for i, name in model.names.items() if name in selected_classes}
-
+    logger.info("Connecting to the IP camera stream...")
     cap = cv2.VideoCapture(stream_url)
-    cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)  # Reduce buffer size
     if not cap.isOpened():
-        return None, "Error: Could not open stream."
+        errors.append("Error: Could not open stream.")
+        return None, "\n".join(errors)
 
-    crossed_objects = {}
-    max_tracked_objects = 1000
+    model = YOLO(model="yolov8n.pt")
+    crossed_objects = set()  # Use a set to store unique object IDs (if available)
 
+    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
 
-        # Optimized inference
-        results = model.track(
-            frame, 
-            persist=True,
-            conf=confidence_threshold,
-            half=True,
-            device=device,
-            verbose=False
-        )
+        # Perform object detection (no tracking)
+        results = model.predict(frame, conf=confidence_threshold)
 
-        if results[0].boxes.id is not None:
-            boxes = results[0].boxes
-            track_ids = boxes.id.int().cpu().tolist()
-            clss = boxes.cls.cpu().tolist()
+        for result in results:
+            boxes = result.boxes.xyxy.cpu().numpy()
+            clss = result.boxes.cls.cpu().numpy()
+            confs = result.boxes.conf.cpu().numpy()
 
-            for box, cls, t_id in zip(boxes.xyxy.cpu(), clss, track_ids):
-                if cls in selected_class_indices and t_id not in crossed_objects:
-                    if is_object_crossing_line(box.numpy(), line_params):
-                        crossed_objects[t_id] = True
-                        if len(crossed_objects) > max_tracked_objects:
-                            crossed_objects.clear()
+            for box, cls, conf in zip(boxes, clss, confs):
+                if conf >= confidence_threshold and model.names[int(cls)] in selected_classes:
+                    # Check if the object crosses the line
+                    if is_object_crossing_line(box, line_params):
+                        # Use the bounding box center as a unique identifier
+                        center = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
+                        crossed_objects.add(tuple(center))  # Add the center to the set
 
-        # Visualization
+        # Visualize the results with bounding boxes
         annotated_frame = results[0].plot()
-        draw_angled_line(annotated_frame, line_params)
-        
-        # Draw count
+
+        # 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
         count = len(crossed_objects)
-        (w, h), _ = cv2.getTextSize(f"COUNT: {count}", cv2.FONT_HERSHEY_SIMPLEX, 1, 2)
-        cv2.rectangle(annotated_frame, (10, 10), (20 + w, 40 + h), (0, 0, 0), -1)
-        cv2.putText(annotated_frame, f"COUNT: {count}", (20, 40), 
-                   cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+        (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.")
 
-# Gradio interface remains unchanged
+# 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>")
+    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")
-    
-    stream_url = gr.Textbox(
-        label="IP Camera Stream URL", 
-        value="https://s104.ipcamlive.com/streams/68idokwtondsqpmkr/stream.m3u8", 
-        visible=False
-    )
-    
-    # First frame extraction
+
+    # 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)
+
+    # Step 1: Extract the first frame from the stream
+    gr.Markdown("### Step 1: Click on the frame to draw a line, the objects crossing it would be counted in real-time.")
     first_frame, status = extract_first_frame(stream_url.value)
-    image = gr.Image(value=first_frame, label="First Frame", type="pil") if first_frame else gr.Markdown(f"**Error:** {status}")
-    line_info = gr.Textbox(label="Line Coordinates", value="Line Coordinates:\nStart: None, End: None")
-    image.select(update_line, inputs=image, outputs=[image, line_info])
-
-    # Class selection
-    class_names = list(model.names.values())
-    selected_classes = gr.CheckboxGroup(choices=class_names, label="Select Classes to Detect")
-
-    # Confidence threshold
-    confidence_threshold = gr.Slider(0.0, 1.0, value=0.2, label="Confidence Threshold")
-
-    # Process button
-    process_button = gr.Button("Process Stream")
-    output_image = gr.Image(label="Processed Frame", streaming=True)
-    error_box = gr.Textbox(label="Errors/Warnings", interactive=False)
-
-    process_button.click(
-        process_video,
-        inputs=[confidence_threshold, selected_classes, stream_url],
-        outputs=[output_image, error_box]
-    )
+    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="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 
+        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(process_video, inputs=[confidence_threshold, selected_classes, stream_url], outputs=[output_image, error_box])
 
+# Launch the interface
 demo.launch(debug=True)