Update app.py

app.py

@@ -8,16 +8,17 @@ from ultralytics import YOLO
 MODEL_PATH = "yolov12x.pt"  # Ensure the model is uploaded to the Hugging Face Space
 model = YOLO(MODEL_PATH)
 
-# COCO dataset class ID for trucks
-TRUCK_CLASS_ID = 7  # "truck"
+# COCO dataset class IDs
+PERSON_CLASS_ID = 0  # "person"
+TRUCK_CLASS_ID = 7    # "truck"
 
-def count_trucks(video_path):
+def count_objects(video_path):
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
         return "Error: Unable to open video file."
 
     frame_count = 0
-    truck_counts = []
+    object_counts = {"people": [], "trucks": []}
     frame_skip = 5  # Process every 5th frame for efficiency
 
     while True:
@@ -32,38 +33,42 @@ def count_trucks(video_path):
         # Run YOLOv12x inference
         results = model(frame, verbose=False)
 
-        truck_count = 0
+        people_count, truck_count = 0, 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 only trucks
-                if class_id == TRUCK_CLASS_ID and confidence > 0.5:
+                # Count objects based on their class IDs
+                if class_id == PERSON_CLASS_ID and confidence > 0.5:
+                    people_count += 1
+                elif class_id == TRUCK_CLASS_ID and confidence > 0.5:
                     truck_count += 1
 
-        truck_counts.append(truck_count)
+        object_counts["people"].append(people_count)
+        object_counts["trucks"].append(truck_count)
 
     cap.release()
 
     return {
-        "Trucks in a Frame": int(np.max(truck_counts)) if truck_counts else 0
+        "Max People in a Frame": int(np.max(object_counts["people"])) if object_counts["people"] else 0,
+        "Max Trucks in a Frame": int(np.max(object_counts["trucks"])) if object_counts["trucks"] else 0
     }
 
 # Gradio UI function
 def analyze_video(video_file):
-    result = count_trucks(video_file)
+    result = count_objects(video_file)
     return "\n".join([f"{key}: {value}" for key, value in result.items()])
 
 # Define Gradio interface
 iface = gr.Interface(
     fn=analyze_video,
     inputs=gr.Video(label="Upload Video"),
-    outputs=gr.Textbox(label="Truck Count"),
-    title="YOLOv12x Truck Counter",
-    description="Upload a video to count trucks using YOLOv12x."
+    outputs=gr.Textbox(label="Analysis Result"),
+    title="YOLOv12x Object Counter",
+    description="Upload a video to count people and trucks using YOLOv12x."
 )
 
 # Launch the Gradio app
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()