Update app.py

app.py

@@ -1,101 +1,86 @@
-# Install required libraries
-#!pip install gradio opencv-python-headless
-
-# Download YOLO files
-#!wget -nc https://raw.githubusercontent.com/pjreddie/darknet/master/cfg/yolov3.cfg
-#!wget -nc https://pjreddie.com/media/files/yolov3.weights
-#!wget -nc https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names
-
 import gradio as gr
 import cv2
 import numpy as np
+import os
+
+# Load YOLO model
+net = cv2.dnn.readNet('yolov3.weights', 'yolov3.cfg')
+
+# Enable GPU (if available)
+net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
+net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
+
+# Load class names
+with open('coco.names', 'r') as f:
+    classes = [line.strip() for line in f.readlines()]
+
+# Get YOLO output layer names
+output_layers_names = net.getUnconnectedOutLayersNames()
 
 def count_people(video_path):
-    # Load YOLO model
-    net = cv2.dnn.readNet('yolov3.weights', 'yolov3.cfg')
-    
-    # Load class names
-    with open('coco.names', 'r') as f:
-        classes = [line.strip() for line in f.readlines()]
-    
-    # Open video
+    if not os.path.exists(video_path):
+        return "Error: Video file not found."
+
     cap = cv2.VideoCapture(video_path)
-    
+    if not cap.isOpened():
+        return "Error: Unable to open video file."
+
     frame_count = 0
-    total_people_count = 0
     people_per_frame = []
-    
-    while cap.isOpened():
+
+    while True:
         ret, frame = cap.read()
         if not ret:
             break
-        
+
         height, width, _ = frame.shape
-        
-        # Create blob from frame
+
+        # Convert frame to YOLO format
         blob = cv2.dnn.blobFromImage(frame, 1/255.0, (416, 416), swapRB=True, crop=False)
         net.setInput(blob)
-        
-        # Get output layer names
-        output_layers_names = net.getUnconnectedOutLayersNames()
-        
+
         # Forward pass
         layer_outputs = net.forward(output_layers_names)
-        
-        # Lists to store detected people
-        boxes = []
-        confidences = []
-        
+
         # Process detections
+        boxes, confidences = [], []
         for output in layer_outputs:
             for detection in output:
                 scores = detection[5:]
                 class_id = np.argmax(scores)
                 confidence = scores[class_id]
-                
-                # Check if detected object is a person
+
                 if classes[class_id] == 'person' and confidence > 0.5:
-                    # Object detected
-                    center_x = int(detection[0] * width)
-                    center_y = int(detection[1] * height)
-                    w = int(detection[2] * width)
-                    h = int(detection[3] * height)
-                    
-                    # Rectangle coordinates
-                    x = int(center_x - w/2)
-                    y = int(center_y - h/2)
-                    
+                    center_x, center_y = int(detection[0] * width), int(detection[1] * height)
+                    w, h = int(detection[2] * width), int(detection[3] * height)
+                    x, y = int(center_x - w / 2), int(center_y - h / 2)
+
                     boxes.append([x, y, w, h])
                     confidences.append(float(confidence))
-        
-        # Apply non-maximum suppression
-        indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
-        
+
+        # Apply Non-Maximum Suppression (NMS)
+        indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4) if boxes else []
+
         # Count people in this frame
-        people_in_frame = len(indexes)
-        people_per_frame.append(people_in_frame)
-        total_people_count += people_in_frame
-        
+        people_per_frame.append(len(indexes))
+
         frame_count += 1
-    
-    # Release resources
+
     cap.release()
-    
-    # Prepare analytics
+
+    # Generate analytics
     return {
-        #'Total Frames Processed': frame_count,
-        #'Total People Detected': total_people_count,
-        #'Average People Per Frame': round(np.mean(people_per_frame), 2),
-        'People in a Video': int(np.max(people_per_frame))  #Max People in a Single Frame
+        "Total Frames Processed": frame_count,
+        "Max People in a Single Frame": int(np.max(people_per_frame)) if people_per_frame else 0,
+        "Avg People Per Frame": round(np.mean(people_per_frame), 2) if people_per_frame else 0
     }
 
-# Define Gradio interface
+# Gradio UI function
 def analyze_video(video_file):
     result = count_people(video_file)
-    result_str = "\n".join([f"{key}: {value}" for key, value in result.items()])
-    return result_str
+    return "\n".join([f"{key}: {value}" for key, value in result.items()])
 
-# Gradio UI
+# Gradio Interface
 interface = gr.Interface(
     fn=analyze_video,
     inputs=gr.Video(label="Upload Video"),
@@ -104,5 +89,5 @@ interface = gr.Interface(
     description="Upload a video to detect and count people using YOLOv3."
 )
 
-# Launch Gradio app
+# Launch app
 interface.launch()