Update app.py

app.py

@@ -2,171 +2,226 @@
 import multiprocessing
 import cv2
 
+# Get the number of CPU cores
+cpu_cores = multiprocessing.cpu_count()
+
 # Set OpenCV to use all available cores
-cv2.setNumThreads(multiprocessing.cpu_count())
+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)
 logger = logging.getLogger(__name__)
 
-# Global variables
-start_point = end_point = line_params = None
+# 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
 
 def extract_first_frame(stream_url):
-    """Extracts first frame from IP camera"""
+    """Extracts the first available frame from the IP camera stream"""
+    logger.info("Extracting first frame...")
     cap = cv2.VideoCapture(stream_url)
     if not cap.isOpened():
         return None, "Error: Could not open stream."
-    
+
     ret, frame = cap.read()
     cap.release()
-    
+
     if not ret:
         return None, "Error: Could not read frame."
-    
-    return Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)), "First frame extracted."
+
+    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(frame_rgb), "First frame extracted successfully."
 
 def update_line(image, evt: gr.SelectData):
     """Handles line drawing interactions"""
     global start_point, end_point, line_params
-    
-    if not start_point:
+
+    if start_point is None:
         start_point = (evt.index[0], evt.index[1])
         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")
-        return image, f"Start: {start_point}"
-    
+        return image, f"Line Coordinates:\nStart: {start_point}, End: None"
+
     end_point = (evt.index[0], evt.index[1])
+    line_params = (start_point, end_point)
+    
     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")
-    
-    # Calculate line parameters
-    if start_point[0] != end_point[0]:
-        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)
-    else:
-        line_params = (float('inf'), start_point[0], start_point, end_point)
-    
+
     start_point = None
-    return image, f"Line: {line_params[2]} to {line_params[3]}"
+    return image, f"Line Coordinates:\nStart: {line_params[0]}, End: {line_params[1]}"
 
-def intersect(A, B, C, D):
-    """Check line segment intersection"""
-    def ccw(A, B, C):
-        return (C[1]-A[1])*(B[0]-A[0]) > (B[1]-A[1])*(C[0]-A[0])
-    
-    return ccw(A,C,D) != ccw(B,C,D) and ccw(A,B,C) != ccw(A,B,D)
+def reset_line():
+    """Resets line coordinates"""
+    global start_point, end_point, line_params
+    start_point = end_point = line_params = None
+    return None, "Line reset. Click to draw a new line."
 
-def is_crossing(box, line_params):
-    """Check if box crosses line"""
+def is_object_crossing_line(box, line_params):
+    """Optimized line crossing check using Liang-Barsky algorithm"""
     if not line_params:
         return False
     
-    (x1, y1), (x2, y2) = line_params[2], line_params[3]
-    box_edges = [
-        ((box[0], box[1]), (box[2], box[1])),
-        ((box[2], box[1]), (box[2], box[3])),
-        ((box[2], box[3]), (box[0], box[3])),
-        ((box[0], box[3]), (box[0], box[1]))
-    ]
-    
-    intersections = 0
-    for edge in box_edges:
-        if intersect((x1,y1), (x2,y2), edge[0], edge[1]):
-            intersections += 1
-            if intersections >= 2:
-                return True
-    return False
-
-def process_video(conf=0.5, classes=None, stream_url=None):
-    """Main processing function"""
+    line_start, line_end = line_params
+    x1, y1, x2, y2 = box
+    return liang_barsky((line_start, line_end), (x1, y1, x2, y2))
+
+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)
+
+def process_video(confidence_threshold=0.5, selected_classes=None, stream_url=None):
+    """Main video processing function with optimizations"""
     global line_params
-    
-    # Initialize YOLOv11
-    model = YOLO('yolo11n.pt')
-    
+    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:
+        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}
+
     cap = cv2.VideoCapture(stream_url)
-    crossed = set()
-    
+    cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)  # Reduce buffer size
+    if not cap.isOpened():
+        return None, "Error: Could not open stream."
+
+    crossed_objects = {}
+    max_tracked_objects = 1000
+
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
-        
-        # Run inference
-        results = model.track(frame, persist=True, conf=conf, classes=classes)
-        
-        # Process results
+
+        # Optimized inference
+        results = model.track(
+            frame, 
+            persist=True,
+            conf=confidence_threshold,
+            half=True,
+            device=device,
+            verbose=False
+        )
+
         if results[0].boxes.id is not None:
-            boxes = results[0].boxes.xyxy.cpu().numpy()
-            ids = results[0].boxes.id.cpu().numpy().astype(int)
-            clss = results[0].boxes.cls.cpu().numpy().astype(int)
-            
-            for box, tid, cls in zip(boxes, ids, clss):
-                if is_crossing(box, line_params) and tid not in crossed:
-                    crossed.add(tid)
-        
-        # Draw overlays
-        annotated = results[0].plot()
-        if line_params:
-            cv2.line(annotated, line_params[2], line_params[3], (0,255,0), 2)
-        cv2.putText(annotated, f"Count: {len(crossed)}", (10,30),
-                   cv2.FONT_HERSHEY_SIMPLEX, 1, (0,255,0), 2)
+            boxes = results[0].boxes
+            track_ids = boxes.id.int().cpu().tolist()
+            clss = boxes.cls.cpu().tolist()
+
+            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()
+
+        # Visualization
+        annotated_frame = results[0].plot()
+        draw_angled_line(annotated_frame, line_params)
         
-        yield cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB), ""
+        # Draw count
+        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)
+
+        yield annotated_frame, ""
 
     cap.release()
 
-# Gradio Interface
-with gr.Blocks() as app:
-    gr.Markdown("# CCTV Object Counter - YOLOv11")
-    
-    # Stream setup
-    url = gr.Textbox(label="Stream URL", value="https://example.com/stream.m3u8")
-    frame_btn = gr.Button("Get First Frame")
-    
-    # Image components
-    img = gr.Image(label="Draw Detection Line", interactive=True)
-    line_info = gr.Textbox(label="Line Coordinates")
-    
-    # Controls
-    classes = gr.CheckboxGroup(label="Classes", choices=[
-        "person", "car", "truck", "motorcycle"
-    ], value=["person"])
-    conf = gr.Slider(0.1, 1.0, value=0.4, label="Confidence Threshold")
-    
-    # Output
-    video_out = gr.Image(label="Live View", streaming=True)
-    status = gr.Textbox(label="Status")
-    
-    # Interactions
-    frame_btn.click(
-        extract_first_frame,
-        inputs=url,
-        outputs=[img, status]
-    )
-    
-    img.select(
-        update_line,
-        inputs=img,
-        outputs=[img, line_info]
+# Gradio interface remains unchanged
+with gr.Blocks() as demo:
+    gr.Markdown("<h1>Real-time monitoring, object tracking, and line-crossing detection for CCTV camera streams.</h1>")
+    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
     )
     
-    gr.Button("Start Counting").click(
+    # First frame extraction
+    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=[conf, classes, url],
-        outputs=[video_out, status]
+        inputs=[confidence_threshold, selected_classes, stream_url],
+        outputs=[output_image, error_box]
     )
 
-app.launch()
+demo.launch(debug=True)