tree3po commited on
Commit
e751200
·
verified ·
1 Parent(s): 8e7d7e0

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +51 -0
app.py ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import gradio as gr
3
+ import supervision as sv
4
+ from ultralytics import YOLO
5
+ from PIL import Image
6
+ import torch
7
+ import time
8
+ import numpy as np
9
+ import uuid
10
+ model = YOLO("yolov8s.pt")
11
+
12
+ def stream_object_detection(video):
13
+ cap = cv2.VideoCapture(video)
14
+ # This means we will output mp4 videos
15
+ video_codec = cv2.VideoWriter_fourcc(*"mp4v") # type: ignore
16
+ fps = int(cap.get(cv2.CAP_PROP_FPS))
17
+ desired_fps = fps // SUBSAMPLE
18
+ width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) // 2
19
+ height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) // 2
20
+ iterating, frame = cap.read()
21
+ n_frames = 0
22
+ output_video_name = f"output_{uuid.uuid4()}.mp4"
23
+ output_video = cv2.VideoWriter(output_video_name, video_codec, desired_fps, (width, height)) # type: ignore
24
+
25
+ while iterating:
26
+ frame = cv2.resize( frame, (0,0), fx=0.5, fy=0.5)
27
+ frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
28
+ result = model(Image.fromarray(frame))[0]
29
+ detections = sv.Detections.from_ultralytics(result)
30
+ outp = draw_box(frame,detections)
31
+ frame = np.array(outp)
32
+ # Convert RGB to BGR
33
+ frame = frame[:, :, ::-1].copy()
34
+ output_video.write(frame)
35
+ batch = []
36
+ output_video.release()
37
+ yield output_video_name
38
+ output_video_name = f"output_{uuid.uuid4()}.mp4"
39
+ output_video = cv2.VideoWriter(output_video_name, video_codec, desired_fps, (width, height)) # type: ignore
40
+ iterating, frame = cap.read()
41
+ n_frames += 1
42
+
43
+ with gr.Blocks() as app:
44
+ #inp = gr.Image(type="filepath")
45
+ with gr.Row():
46
+ with gr.Column():
47
+ inp = gr.Video()
48
+ btn = gr.Button()
49
+ outp_v = gr.Video(label="Processed Video", streaming=True, autoplay=True)
50
+ btn.click(stream_object_detection,inp,[outp_v])
51
+ app.queue(concurrency_limit=20).launch()