Update app.py

app.py

@@ -1,6 +1,4 @@
-##############
-
-#Maximize CPU usage
+# Maximize CPU usage and GPU utilization
 import multiprocessing
 import cv2
 
@@ -14,7 +12,7 @@ cv2.setNumThreads(cpu_cores)
 print(f"OpenCV using {cv2.getNumThreads()} threads out of {cpu_cores} available cores")
 
 ##############
-import cv2
+import torch
 import gradio as gr
 import numpy as np
 from PIL import Image, ImageDraw
@@ -32,284 +30,158 @@ 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
+
 def extract_first_frame(stream_url):
-    """
-    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...")
+    """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("Error: Could not open stream.")
+        logger.error("Could not open stream.")
         return None, "Error: Could not open stream."
 
     ret, frame = cap.read()
     cap.release()
 
     if not ret:
-        logger.error("Error: Could not read the first frame.")
-        return None, "Error: Could not read the first frame."
+        logger.error("Could not read frame.")
+        return None, "Error: Could not read frame."
 
-    # Convert the frame to a PIL image
     frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-    pil_image = Image.fromarray(frame_rgb)
-
-    logger.info("First frame extracted successfully.")
-    return pil_image, "First frame extracted successfully."
+    return Image.fromarray(frame_rgb), "First frame extracted."
 
 def update_line(image, evt: gr.SelectData):
-    """
-    Updates the line based on user interaction (click and drag).
-    """
+    """Updates the line based on user interaction."""
     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 Start: {start_point}"
 
-        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])
+    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 the slope (m) and intercept (b) of the line: y = mx + b
-    if start_point[0] != end_point[0]:  # Avoid division by zero
+    # 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)  # Store slope, intercept, and points
+        line_params = (slope, intercept, start_point, end_point)
     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"
-    )
-
-    # 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
+    return image, f"Line: {line_params[0]:.2f}x + {line_params[1]:.2f}"
 
-def reset_line():
-    """
-    Resets the line coordinates.
-    """
-    global start_point, end_point, line_params
-    start_point = None
-    end_point = None
-    line_params = None
-    return None, "Line reset. Click to draw a new line."
-
-def intersect(A, B, C, D):
-    """
-    Determines if two line segments AB and CD intersect.
-    """
-    def ccw(A, B, C):
-        return (C[1] - A[1]) * (B[0] - A[0]) - (B[1] - A[1]) * (C[0] - A[0])
-
-    def on_segment(A, B, C):
-        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)
-
-    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):
-    """
-    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
-
-    # 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_point, end_point = line_params
-    cv2.line(image, start_point, end_point, color, thickness)
+def optimized_intersection_check(box, line_params):
+    """Optimized line-box intersection check using vector math."""
+    _, _, (x1, y1), (x2, y2) = line_params
+    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')
+    
+    for i in range(2):
+        if dx == 0 and dy == 0:
+            continue
+            
+        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
+            
+        t_min = min(t0, t1)
+        t_max = max(t0, t1)
+        
+        if t_min > t_near: t_near = t_min
+        if t_max < t_far: t_far = t_max
+        
+    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):
-    """
-    Processes the IP camera stream to count objects of the selected classes crossing the line.
-    """
+    """Optimized video processing pipeline."""
     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.")
+    # Validation checks
+    if not line_params or not selected_classes or not stream_url:
+        return None, "Missing configuration parameters"
 
-    if errors:
-        return None, "\n".join(errors)
-
-    logger.info("Connecting to the IP camera stream...")
+    # Convert to set for faster lookups
+    selected_classes = set(selected_classes)
+    
+    # Video capture setup
     cap = cv2.VideoCapture(stream_url)
     if not cap.isOpened():
-        errors.append("Error: Could not open stream.")
-        return None, "\n".join(errors)
-
-    model = YOLO(model="yolo11n.pt")
-    crossed_objects = {}
-    max_tracked_objects = 1000  # Maximum number of objects to track before clearing
+        return None, "Error opening stream"
+    
+    crossed_objects = set()
+    frame_count = 0
 
-    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
-
-        # Perform object tracking with confidence threshold
-        results = model.track(frame, persist=True, conf=confidence_threshold)
-
+            
+        frame_count += 1
+        if frame_count % FRAME_SKIP != 0:
+            continue
+            
+        # Preprocess frame
+        frame = cv2.resize(frame, None, fx=FRAME_SCALE, fy=FRAME_SCALE)
+        
+        # Object detection
+        results = model.track(
+            frame,
+            persist=True,
+            conf=confidence_threshold,
+            verbose=False,
+            device=device,
+            tracker="botsort.yaml"  # Use optimized tracker config
+        )
+        
+        # Process detections
         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
+            boxes = results[0].boxes.xyxy.cpu().numpy()
+            track_ids = results[0].boxes.id.int().cpu().numpy()
+            classes = 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()
+
+        # Annotation
         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
+        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)
+        
         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
-    # gr.Markdown("### Step 0: Enter the IP Camera Stream URL")
-    # stream_url = gr.Textbox(label="Enter IP Camera Stream URL", value="https://s103.ipcamlive.com/streams/67n4ojknye7lkxpmf/stream.m3u8", visible=False)
-    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)
-    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])
-
-        # Reset the line (optional)
-        # gr.Markdown("### Step 4: Reset the Line (Optional)")
-        # reset_button = gr.Button("Reset Line")
-        # reset_button.click(reset_line, inputs=None, outputs=[image, line_info])
-
-        # Step 2: Select classes to detect
-        gr.Markdown("### Step 2: Select Classes to Detect")
-        model = YOLO(model="yolo11n.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)