Update app.py

app.py

@@ -1,15 +1,9 @@
-# Maximize CPU usage and GPU utilization
+# Maximize performance settings
 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")
+# Configure OpenCV for multi-core processing
+cv2.setNumThreads(multiprocessing.cpu_count())
 
 ##############
 import torch
@@ -17,171 +11,220 @@ 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 time
 
-# Set up logging
+# Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Global variables to store line coordinates and line equation
-start_point = None
-end_point = None
-line_params = None  # Stores (slope, intercept) of the line
-
-# Initialize model once
-model = YOLO('yolov8n.pt')  # Use smaller model if needed
-# Check for GPU availability
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-model.to(device)
-logger.info(f"Using device: {device}")
-
-# Video processing parameters
-FRAME_SKIP = 1  # Process every nth frame
-FRAME_SCALE = 0.5  # Scale factor for input frames
+# Global variables for line coordinates
+line_params = None
+model = None
+
+def initialize_yolov11():
+    """Initialize YOLOv11 model with error handling"""
+    global model
+    try:
+        model = YOLO('yolov11n.pt')  # Make sure this model file exists
+        if torch.cuda.is_available():
+            model.to('cuda')
+            logger.info("YOLOv11 initialized with CUDA acceleration")
+        else:
+            logger.info("YOLOv11 initialized with CPU")
+        return True
+    except Exception as e:
+        logger.error(f"Model initialization failed: {str(e)}")
+        return False
 
 def extract_first_frame(stream_url):
-    """Extracts the first available frame from the IP camera stream."""
-    logger.info("Extracting first frame...")
-    cap = cv2.VideoCapture(stream_url)
-    if not cap.isOpened():
-        logger.error("Could not open stream.")
-        return None, "Error: Could not open stream."
-
-    ret, frame = cap.read()
-    cap.release()
-
-    if not ret:
-        logger.error("Could not read frame.")
-        return None, "Error: Could not read frame."
-
-    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-    return Image.fromarray(frame_rgb), "First frame extracted."
+    """Robust frame extraction with retries"""
+    for _ in range(3):  # Retry up to 3 times
+        cap = cv2.VideoCapture(stream_url)
+        if cap.isOpened():
+            ret, frame = cap.read()
+            cap.release()
+            if ret:
+                return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), "First frame extracted"
+            time.sleep(1)  # Wait before retry
+    return None, "Error: Failed to capture initial frame"
 
 def update_line(image, evt: gr.SelectData):
-    """Updates the line based on user interaction."""
-    global start_point, end_point, line_params
-
-    if start_point is None:
-        start_point = (evt.index[0], evt.index[1])
+    """Optimized line drawing with validation"""
+    global line_params
+    
+    if not hasattr(image, 'points'):
+        image.points = []
+    
+    if len(image.points) < 2:
+        image.points.append((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"Line Start: {start_point}"
-
-    end_point = (evt.index[0], evt.index[1])
-    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[0]:.2f}x + {line_params[1]:.2f}"
+        color = "blue" if len(image.points) == 1 else "green"
+        draw.ellipse([evt.index[0]-5, evt.index[1]-5, evt.index[0]+5, evt.index[1]+5], 
+                    fill=color, outline=color)
+        
+    if len(image.points) == 2:
+        x1, y1 = image.points[0]
+        x2, y2 = image.points[1]
+        draw = ImageDraw.Draw(image)
+        draw.line([(x1,y1), (x2,y2)], fill="red", width=2)
+        
+        # Store line parameters
+        if x2 - x1 != 0:
+            slope = (y2 - y1) / (x2 - x1)
+            intercept = y1 - slope * x1
+        else:
+            slope = float('inf')
+            intercept = x1
+        line_params = (slope, intercept, (x1,y1), (x2,y2))
+        
+    status = f"Points: {len(image.points)}/2" if len(image.points) < 2 else "Line set!"
+    return image, status
 
-def optimized_intersection_check(box, line_params):
-    """Optimized line-box intersection check using vector math."""
-    _, _, (x1, y1), (x2, y2) = line_params
+def line_intersection(box, line):
+    """Fast line-box intersection using vector math"""
+    (m, b, (x1,y1), (x2,y2)) = line
     box_x1, box_y1, box_x2, box_y2 = box
     
     # Convert line to parametric form
     dx = x2 - x1
     dy = y2 - y1
     
-    # Check if any box edge intersects the line
-    t_near = -float('inf')
-    t_far = float('inf')
+    # Check box edges
+    t0 = 0.0
+    t1 = 1.0
     
-    for i in range(2):
-        if dx == 0 and dy == 0:
-            continue
+    for edge in [0, 1]:  # Check both x and y axes
+        if edge == 0:  # X-axis boundaries
+            dir = dx
+            p = box_x1 - x1
+            q = box_x2 - x1
+        else:  # Y-axis boundaries
+            dir = dy
+            p = box_y1 - y1
+            q = box_y2 - y1
             
-        if i == 0:  # X-axis
-            t0 = (box_x1 - x1) / dx if dx != 0 else 0
-            t1 = (box_x2 - x1) / dx if dx != 0 else 0
-        else:  # Y-axis
-            t0 = (box_y1 - y1) / dy if dy != 0 else 0
-            t1 = (box_y2 - y1) / dy if dy != 0 else 0
+        if dir == 0:
+            if p > 0 or q < 0: return False
+            continue
             
-        t_min = min(t0, t1)
-        t_max = max(t0, t1)
+        t_near = p / dir
+        t_far = q / dir
+        if t_near > t_far: t_near, t_far = t_far, t_near
         
-        if t_min > t_near: t_near = t_min
-        if t_max < t_far: t_far = t_max
+        t0 = max(t0, t_near)
+        t1 = min(t1, t_far)
         
-    return t_near <= t_far and t_near <= 1 and t_far >= 0
-
-def process_video(confidence_threshold=0.5, selected_classes=None, stream_url=None):
-    """Optimized video processing pipeline."""
-    global line_params
-
-    # Validation checks
-    if not line_params or not selected_classes or not stream_url:
-        return None, "Missing configuration parameters"
+        if t0 > t1: return False
+    
+    return t0 <= 1 and t1 >= 0
 
-    # Convert to set for faster lookups
-    selected_classes = set(selected_classes)
+def process_stream(conf_thresh, classes, stream_url):
+    """Optimized video processing pipeline"""
+    if not model:
+        yield None, "Model not initialized"
+        return
+    
+    if not line_params:
+        yield None, "No detection line set"
+        return
     
-    # Video capture setup
     cap = cv2.VideoCapture(stream_url)
     if not cap.isOpened():
-        return None, "Error opening stream"
+        yield None, "Failed to open video stream"
+        return
     
-    crossed_objects = set()
-    frame_count = 0
-
-    while cap.isOpened():
+    tracker = {}  # {track_id: last_seen}
+    crossed = set()
+    frame_skip = 2  # Process every 2nd frame
+    count = 0
+    
+    while True:
         ret, frame = cap.read()
         if not ret:
             break
-            
-        frame_count += 1
-        if frame_count % FRAME_SKIP != 0:
+        
+        count += 1
+        if count % frame_skip != 0:
             continue
-            
-        # Preprocess frame
-        frame = cv2.resize(frame, None, fx=FRAME_SCALE, fy=FRAME_SCALE)
         
-        # Object detection
+        # Detection
         results = model.track(
             frame,
             persist=True,
-            conf=confidence_threshold,
+            conf=conf_thresh,
+            classes=classes,
             verbose=False,
-            device=device,
-            tracker="botsort.yaml"  # Use optimized tracker config
+            device='cuda' if torch.cuda.is_available() else 'cpu'
         )
         
-        # Process detections
+        # Processing
         if results[0].boxes.id is not None:
             boxes = results[0].boxes.xyxy.cpu().numpy()
-            track_ids = results[0].boxes.id.int().cpu().numpy()
-            classes = results[0].boxes.cls.cpu().numpy()
+            ids = results[0].boxes.id.int().cpu().numpy()
+            scores = results[0].boxes.conf.cpu().numpy()
+            labels = results[0].boxes.cls.cpu().numpy()
             
-            for box, track_id, cls in zip(boxes, track_ids, classes):
-                if model.names[int(cls)] not in selected_classes:
-                    continue
-                
-                if optimized_intersection_check(box, line_params) and track_id not in crossed_objects:
-                    crossed_objects.add(track_id)
-                    if len(crossed_objects) > 1000:
-                        crossed_objects.clear()
-
+            for box, track_id, score, label in zip(boxes, ids, scores, labels):
+                if line_intersection(box, line_params) and track_id not in crossed:
+                    crossed.add(track_id)
+                    if len(crossed) > 1000:
+                        crossed.clear()
+        
         # Annotation
-        annotated_frame = results[0].plot()
-        cv2.line(annotated_frame, line_params[2], line_params[3], (0,255,0), 2)
-        cv2.putText(annotated_frame, f"COUNT: {len(crossed_objects)}", 
-                   (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,255,0), 2)
+        annotated = results[0].plot()
+        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)
         
-        yield annotated_frame, ""
+        yield cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB), ""
 
     cap.release()
 
+# Gradio Interface
+with gr.Blocks() as app:
+    gr.Markdown("# CCTV Smart Monitor - YOLOv11")
+    
+    # Initialization
+    if not initialize_yolov11():
+        gr.Markdown("**Error**: Failed to initialize YOLOv11 model")
+    
+    # Stream URL input
+    stream_url = gr.Textbox(
+        label="RTSP Stream URL",
+        value="rtsp://example.com/stream",
+        visible=True
+    )
+    
+    # Frame setup
+    with gr.Row():
+        frame = gr.Image(label="Setup Frame", interactive=True)
+        line_status = gr.Textbox(label="Line Status", interactive=False)
+    
+    # Controls
+    with gr.Row():
+        class_selector = gr.CheckboxGroup(
+            choices=model.names.values() if model else [],
+            label="Detection Classes"
+        )
+        confidence = gr.Slider(0.1, 1.0, value=0.4, label="Confidence Threshold")
+    
+    # Output
+    output_video = gr.Image(label="Live Analysis", streaming=True)
+    error_box = gr.Textbox(label="System Messages", interactive=False)
+    
+    # Interactions
+    frame.select(
+        update_line,
+        inputs=frame,
+        outputs=[frame, line_status]
+    )
+    
+    gr.Button("Start Analysis").click(
+        process_stream,
+        inputs=[confidence, class_selector, stream_url],
+        outputs=[output_video, error_box]
+    )
+
+if __name__ == "__main__":
+    app.launch(debug=True, enable_queue=True)