hb-setosys commited on
Commit
0a9b429
·
verified ·
1 Parent(s): a027dbb

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +74 -0
app.py ADDED
@@ -0,0 +1,74 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import numpy as np
3
+ import torch
4
+ import gradio as gr
5
+ from ultralytics import YOLO
6
+
7
+ # Load YOLOv12x model
8
+ MODEL_PATH = "yolov12x.pt" # Ensure the model is uploaded to the Hugging Face Space
9
+ model = YOLO(MODEL_PATH)
10
+
11
+ # COCO dataset class IDs
12
+ PERSON_CLASS_ID = 0 # "person"
13
+ TRUCK_CLASS_ID = 7 # "truck"
14
+
15
+ def count_objects(video_path):
16
+ cap = cv2.VideoCapture(video_path)
17
+ if not cap.isOpened():
18
+ return "Error: Unable to open video file."
19
+
20
+ frame_count = 0
21
+ object_counts = {"people": [], "trucks": []}
22
+ frame_skip = 5 # Process every 5th frame for efficiency
23
+
24
+ while True:
25
+ ret, frame = cap.read()
26
+ if not ret:
27
+ break # End of video
28
+
29
+ frame_count += 1
30
+ if frame_count % frame_skip != 0:
31
+ continue # Skip frames to improve efficiency
32
+
33
+ # Run YOLOv12x inference
34
+ results = model(frame, verbose=False)
35
+
36
+ people_count, truck_count = 0, 0
37
+ for result in results:
38
+ for box in result.boxes:
39
+ class_id = int(box.cls.item()) # Get class ID
40
+ confidence = float(box.conf.item()) # Get confidence score
41
+
42
+ # Count objects based on their class IDs
43
+ if class_id == PERSON_CLASS_ID and confidence > 0.5:
44
+ people_count += 1
45
+ elif class_id == TRUCK_CLASS_ID and confidence > 0.5:
46
+ truck_count += 1
47
+
48
+ object_counts["people"].append(people_count)
49
+ object_counts["trucks"].append(truck_count)
50
+
51
+ cap.release()
52
+
53
+ return {
54
+ "Max People in a Frame": int(np.max(object_counts["people"])) if object_counts["people"] else 0,
55
+ "Max Trucks in a Frame": int(np.max(object_counts["trucks"])) if object_counts["trucks"] else 0
56
+ }
57
+
58
+ # Gradio UI function
59
+ def analyze_video(video_file):
60
+ result = count_objects(video_file)
61
+ return "\n".join([f"{key}: {value}" for key, value in result.items()])
62
+
63
+ # Define Gradio interface
64
+ iface = gr.Interface(
65
+ fn=analyze_video,
66
+ inputs=gr.Video(label="Upload Video"),
67
+ outputs=gr.Textbox(label="Analysis Result"),
68
+ title="YOLOv12x Object Counter",
69
+ description="Upload a video to count people and trucks using YOLOv12x."
70
+ )
71
+
72
+ # Launch the Gradio app
73
+ if __name__ == "__main__":
74
+ iface.launch()