Update app.py

app.py

@@ -1,179 +1,126 @@
-# Maximize performance settings
+# Maximize CPU usage
 import multiprocessing
 import cv2
 
-# Configure OpenCV for multi-core processing
+# Set OpenCV to use all available cores
 cv2.setNumThreads(multiprocessing.cpu_count())
 
 ##############
-import torch
 import gradio as gr
 import numpy as np
 from PIL import Image, ImageDraw
 from ultralytics import YOLO
 import logging
-import time
 
-# Configure logging
+# Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# 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
+# Global variables
+start_point = end_point = line_params = None
 
 def extract_first_frame(stream_url):
-    """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"
+    """Extracts first frame from IP camera"""
+    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."
 
 def update_line(image, evt: gr.SelectData):
-    """Optimized line drawing with validation"""
-    global line_params
-    
-    if not hasattr(image, 'points'):
-        image.points = []
+    """Handles line drawing interactions"""
+    global start_point, end_point, line_params
     
-    if len(image.points) < 2:
-        image.points.append((evt.index[0], evt.index[1]))
-        draw = ImageDraw.Draw(image)
-        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]
+    if not start_point:
+        start_point = (evt.index[0], evt.index[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
+        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}"
+    
+    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[2]} to {line_params[3]}"
 
-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 box edges
-    t0 = 0.0
-    t1 = 1.0
-    
-    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 dir == 0:
-            if p > 0 or q < 0: return False
-            continue
-            
-        t_near = p / dir
-        t_far = q / dir
-        if t_near > t_far: t_near, t_far = t_far, t_near
-        
-        t0 = max(t0, t_near)
-        t1 = min(t1, t_far)
-        
-        if t0 > t1: return False
+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 t0 <= 1 and t1 >= 0
+    return ccw(A,C,D) != ccw(B,C,D) and ccw(A,B,C) != ccw(A,B,D)
 
-def process_stream(conf_thresh, classes, stream_url):
-    """Optimized video processing pipeline"""
-    if not model:
-        yield None, "Model not initialized"
-        return
-    
+def is_crossing(box, line_params):
+    """Check if box crosses line"""
     if not line_params:
-        yield None, "No detection line set"
-        return
+        return False
     
-    cap = cv2.VideoCapture(stream_url)
-    if not cap.isOpened():
-        yield None, "Failed to open video stream"
-        return
+    (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"""
+    global line_params
     
-    tracker = {}  # {track_id: last_seen}
+    # Initialize YOLOv11
+    model = YOLO('yolo11n.pt')
+    
+    cap = cv2.VideoCapture(stream_url)
     crossed = set()
-    frame_skip = 2  # Process every 2nd frame
-    count = 0
     
-    while True:
+    while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
         
-        count += 1
-        if count % frame_skip != 0:
-            continue
+        # Run inference
+        results = model.track(frame, persist=True, conf=conf, classes=classes)
         
-        # Detection
-        results = model.track(
-            frame,
-            persist=True,
-            conf=conf_thresh,
-            classes=classes,
-            verbose=False,
-            device='cuda' if torch.cuda.is_available() else 'cpu'
-        )
-        
-        # Processing
+        # Process results
         if results[0].boxes.id is not None:
             boxes = results[0].boxes.xyxy.cpu().numpy()
-            ids = results[0].boxes.id.int().cpu().numpy()
-            scores = results[0].boxes.conf.cpu().numpy()
-            labels = results[0].boxes.cls.cpu().numpy()
+            ids = results[0].boxes.id.cpu().numpy().astype(int)
+            clss = results[0].boxes.cls.cpu().numpy().astype(int)
             
-            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()
+            for box, tid, cls in zip(boxes, ids, clss):
+                if is_crossing(box, line_params) and tid not in crossed:
+                    crossed.add(tid)
         
-        # Annotation
+        # Draw overlays
         annotated = results[0].plot()
-        cv2.line(annotated, line_params[2], line_params[3], (0,255,0), 2)
+        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)
         
@@ -183,48 +130,43 @@ def process_stream(conf_thresh, classes, stream_url):
 
 # Gradio Interface
 with gr.Blocks() as app:
-    gr.Markdown("# CCTV Smart Monitor - YOLOv11")
+    gr.Markdown("# CCTV Object Counter - YOLOv11")
     
-    # Initialization
-    if not initialize_yolov11():
-        gr.Markdown("**Error**: Failed to initialize YOLOv11 model")
+    # Stream setup
+    url = gr.Textbox(label="Stream URL", value="https://example.com/stream.m3u8")
+    frame_btn = gr.Button("Get First Frame")
     
-    # 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)
+    # Image components
+    img = gr.Image(label="Draw Detection Line", interactive=True)
+    line_info = gr.Textbox(label="Line Coordinates")
     
     # 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")
+    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
-    output_video = gr.Image(label="Live Analysis", streaming=True)
-    error_box = gr.Textbox(label="System Messages", interactive=False)
+    video_out = gr.Image(label="Live View", streaming=True)
+    status = gr.Textbox(label="Status")
     
     # Interactions
-    frame.select(
+    frame_btn.click(
+        extract_first_frame,
+        inputs=url,
+        outputs=[img, status]
+    )
+    
+    img.select(
         update_line,
-        inputs=frame,
-        outputs=[frame, line_status]
+        inputs=img,
+        outputs=[img, line_info]
     )
     
-    gr.Button("Start Analysis").click(
-        process_stream,
-        inputs=[confidence, class_selector, stream_url],
-        outputs=[output_video, error_box]
+    gr.Button("Start Counting").click(
+        process_video,
+        inputs=[conf, classes, url],
+        outputs=[video_out, status]
     )
 
-if __name__ == "__main__":
-    app.launch(debug=True, enable_queue=True)
+app.launch()