Update app.py

app.py

@@ -1,96 +1,66 @@
 import gradio as gr
 import cv2
 import numpy as np
-import os
+import torch
+from ultralytics import YOLO
 
-# Load YOLO model
-net = cv2.dnn.readNet('yolov3.weights', 'yolov3.cfg')
+# Load YOLOv8 model (pre-trained on COCO dataset)
+model = YOLO("yolov8x.pt")  # Use "yolov8x.pt" for highest accuracy
 
-# Enable GPU (if available)
-#net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
-#net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
-
-# Set backend to CPU
-net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
-net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
-
-# 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):
-    if not os.path.exists(video_path):
-        return "Error: Video file not found."
+# Class label for trucks (COCO dataset)
+TRUCK_CLASS_ID = 7  # "truck" in COCO dataset
 
+def count_trucks(video_path):
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
         return "Error: Unable to open video file."
 
     frame_count = 0
-    people_per_frame = []
+    truck_count_per_frame = []
+    frame_skip = 5  # Process every 5th frame for efficiency
 
     while True:
         ret, frame = cap.read()
         if not ret:
-            break
-
-        height, width, _ = frame.shape
+            break  # End of video
 
-        # Convert frame to YOLO format
-        blob = cv2.dnn.blobFromImage(frame, 1/255.0, (416, 416), swapRB=True, crop=False)
-        net.setInput(blob)
-
-        # Forward pass
-        layer_outputs = net.forward(output_layers_names)
-
-        # 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]
-
-                if classes[class_id] == 'person' and confidence > 0.5:
-                    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)
+        frame_count += 1
+        if frame_count % frame_skip != 0:
+            continue  # Skip frames to improve efficiency
 
-                    boxes.append([x, y, w, h])
-                    confidences.append(float(confidence))
+        # Run YOLOv8 inference
+        results = model(frame, verbose=False)
 
-        # Apply Non-Maximum Suppression (NMS)
-        indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4) if boxes else []
+        truck_count = 0
+        for result in results:
+            for box in result.boxes:
+                class_id = int(box.cls.item())  # Get class ID
+                confidence = float(box.conf.item())  # Get confidence score
 
-        # Count people in this frame
-        people_per_frame.append(len(indexes))
+                if class_id == TRUCK_CLASS_ID and confidence > 0.6:
+                    truck_count += 1  # Count trucks with confidence > 0.6
 
-        frame_count += 1
+        truck_count_per_frame.append(truck_count)
 
     cap.release()
 
-    # Generate analytics
     return {
-        #"Total Frames Processed": frame_count,
-        "People in a Video": 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
+        "Total Trucks in Video": int(np.max(truck_count_per_frame)) if truck_count_per_frame else 0,
+        "Avg Trucks Per Frame": round(np.mean(truck_count_per_frame), 2) if truck_count_per_frame else 0
     }
 
 # Gradio UI function
 def analyze_video(video_file):
-    result = count_people(video_file)
+    result = count_trucks(video_file)
     return "\n".join([f"{key}: {value}" for key, value in result.items()])
 
 # Gradio Interface
 interface = gr.Interface(
     fn=analyze_video,
     inputs=gr.Video(label="Upload Video"),
-    outputs=gr.Textbox(label="People Counting Results"),
-    title="YOLO-based People Counter",
-    description="Upload a video to detect and count people using YOLOv3."
+    outputs=gr.Textbox(label="Truck Counting Results"),
+    title="YOLOv8-based Truck Counter",
+    description="Upload a video to detect and count trucks using YOLOv8."
 )
 
 # Launch app