am-om commited on
Commit
2cce14f
·
1 Parent(s): aa6397f
Files changed (2) hide show
  1. app.py +124 -0
  2. requirements.txt +0 -0
app.py ADDED
@@ -0,0 +1,124 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+
3
+
4
+
5
+ import os
6
+ import json
7
+ from ultralytics import YOLO
8
+ import supervision as sv
9
+
10
+
11
+ # --- 1. CONFIGURATION ---
12
+
13
+ # Folder where the script is
14
+ BASE_DIR = os.path.dirname(os.path.abspath(__file__))
15
+
16
+ # Go up one level to Om_Singh
17
+ PROJECT_DIR = os.path.dirname(BASE_DIR) # parent folder of app/
18
+
19
+ # Paths for outputs in video_and_json folder
20
+ VIDEO_DIR = os.path.join(PROJECT_DIR, "video_and_json") # Om_Singh/video_and_json
21
+ os.makedirs(VIDEO_DIR, exist_ok=True) # ensure folder exists
22
+
23
+ MODEL_PATH = os.path.join(BASE_DIR, "best.pt") # model stays in app/
24
+ OUTPUT_VIDEO_PATH = os.path.join(VIDEO_DIR, "output.mp4")
25
+ OUTPUT_JSON_PATH = os.path.join(VIDEO_DIR, "result.json")
26
+
27
+ def process_video(INPUT_VIDEO_PATH, OUTPUT_VIDEO_PATH, OUTPUT_JSON_PATH):
28
+ print("Loading model...")
29
+ model = YOLO(MODEL_PATH)
30
+
31
+ print("Initializing tracker and annotators...")
32
+ tracker = sv.ByteTrack()
33
+ box_annotator = sv.BoxAnnotator(thickness=5)
34
+ label_annotator = sv.LabelAnnotator(text_position=sv.Position.TOP_CENTER, text_scale = 1, text_thickness= 1)
35
+
36
+ sv.LabelAnnotator()
37
+
38
+ frame_generator = sv.get_video_frames_generator(source_path=INPUT_VIDEO_PATH)
39
+ video_info = sv.VideoInfo.from_video_path(INPUT_VIDEO_PATH)
40
+
41
+ results_list = []
42
+
43
+ with sv.VideoSink(target_path=OUTPUT_VIDEO_PATH, video_info=video_info) as sink:
44
+ print("Processing video frames...")
45
+ for frame_number, frame in enumerate(frame_generator):
46
+ # Run YOLO prediction
47
+ results = model(frame)[0]
48
+ detections = sv.Detections.from_ultralytics(results)
49
+
50
+ # Update tracker
51
+ tracked_detections = tracker.update_with_detections(detections=detections)
52
+
53
+ # Prepare labels
54
+ labels = [
55
+ f"ID: {det[4]} {model.model.names[int(det[3])]}"
56
+ for det in tracked_detections
57
+ ]
58
+
59
+ # Annotate frame
60
+ annotated_frame = box_annotator.annotate(scene=frame.copy(), detections=tracked_detections)
61
+ annotated_frame = label_annotator.annotate(scene=annotated_frame, detections=tracked_detections, labels=labels)
62
+
63
+ # Save tracking info
64
+ for det in tracked_detections:
65
+ bbox = det[0]
66
+ conf = det[2]
67
+ class_id = int(det[3])
68
+ tracker_id = det[4]
69
+
70
+ results_list.append({
71
+ "frame_number": frame_number,
72
+ "track_id": int(tracker_id),
73
+ "class": model.model.names[class_id],
74
+ "confidence": float(conf),
75
+ "bounding_box": [int(coord) for coord in bbox]
76
+ })
77
+
78
+ # Write annotated frame
79
+ sink.write_frame(frame=annotated_frame)
80
+
81
+ if frame_number % 30 == 0:
82
+ print(f"Processed frame {frame_number}...")
83
+
84
+ print("Video processing complete. Saving results...")
85
+ with open(OUTPUT_JSON_PATH, 'w') as f:
86
+ json.dump(results_list, f, indent=4)
87
+
88
+ print("--- All tasks finished successfully! ---")
89
+
90
+
91
+
92
+ # --- Main processing function ---
93
+ def process(input_video):
94
+ output_video = "output.mp4"
95
+ output_json = "result.json"
96
+
97
+ # During processing: red text
98
+ status_html = "<p style='color:red; font-weight:bold;'>Processing...</p>"
99
+
100
+ # Run video processing
101
+ process_video(input_video, output_video, output_json)
102
+
103
+ # After processing: green text
104
+ status_html = "<p style='color:limegreen; font-weight:bold;'>Processing complete!</p>"
105
+ return status_html, output_video, output_json
106
+
107
+ # --- Gradio UI ---
108
+ with gr.Blocks() as demo:
109
+ gr.Markdown("<h1 style='text-align:center;'>Vehicle and Pedestrian Tracker</h1>")
110
+
111
+ input_video = gr.Video(label="Upload Video")
112
+ start_btn = gr.Button("Start Tracking")
113
+ status_display = gr.HTML("") # Initially empty
114
+ output_video = gr.Video(label="Processed Video")
115
+ output_json = gr.File(label="Download JSON Output")
116
+
117
+ start_btn.click(
118
+ fn=process,
119
+ inputs=input_video,
120
+ outputs=[status_display, output_video, output_json]
121
+ )
122
+
123
+ if __name__ == "__main__":
124
+ demo.launch(share = True)
requirements.txt ADDED
Binary file (4.16 kB). View file