Done
Browse files- app.py +124 -0
- 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
|
|
|