Spaces:
Running
Running
| import gradio as gr | |
| import time | |
| import cv2 # opencv2 package for python. | |
| import torch | |
| from pytube import YouTube | |
| from ultralyticsplus import YOLO, render_result | |
| model = YOLO('ultralyticsplus/yolov8s') | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| URL = "https://www.youtube.com/watch?v=6NBwbKMyzEE" #URL to parse | |
| # set model parameters | |
| model.overrides['conf'] = 0.50 # NMS confidence threshold | |
| model.overrides['iou'] = 0.45 # NMS IoU threshold | |
| model.overrides['agnostic_nms'] = False # NMS class-agnostic | |
| model.overrides['max_det'] = 1000 # maximum number of detections per image | |
| model.to(device) | |
| def load(URL): | |
| yt = YouTube(URL) | |
| vid_cap = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').desc().last().download(filename="tmp.mp4") | |
| global player | |
| player = cv2.VideoCapture(vid_cap) | |
| frame_num = int(player.get(cv2.CAP_PROP_POS_FRAMES)) | |
| frame_count = int(player.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| frame_fps = (player.get(cv2.CAP_PROP_FPS)) | |
| tog = 0 | |
| return vid_cap,frame_num,frame_count,frame_fps,tog | |
| def vid_play(cap,frame_num): | |
| assert player.isOpened() # Make sure that their is a stream. | |
| player.set(cv2.CAP_PROP_POS_FRAMES, int(frame_num)) | |
| ret, frame_bgr = player.read(int(frame_num)) | |
| frame = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) | |
| results = model.predict(frame) | |
| render = render_result(model=model, image=frame, result=results[0]) | |
| return render | |
| def fw_fn(cur,last): | |
| next = cur+1 | |
| if next > last: | |
| next = last | |
| return next | |
| def bk_fn(cur): | |
| next = cur-1 | |
| if next < 0: | |
| next = 0 | |
| return next | |
| def tog_on(): | |
| return 1,gr.Markdown.update("""<center><h7>Status: Playing 😁</h7></center>""") | |
| def tog_off(): | |
| return 0,gr.Markdown.update("""<center><h7>Status: Stopped 💀</h7></center>""") | |
| def pl_fn(cap,cur,last,fps,pl_tog): | |
| player.set(cv2.CAP_PROP_POS_FRAMES, cur) | |
| ret, frame_bgr = player.read(cur) | |
| frame = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) | |
| results = model.predict(frame) | |
| render = render_result(model=model, image=frame, result=results[0]) | |
| if pl_tog ==1: | |
| cur+=1 | |
| else: | |
| cur = cur | |
| return render,cur | |
| with gr.Blocks() as app: | |
| gr.Markdown("""<center><h1>Slow Video Object Detection</h1><h4>Gradio and ultralyticsplus/yolov8s</h4><h4>Probably faster on GPU 🤷♂️</h4></center>""") | |
| play_state = gr.Markdown("""<right><h7></h7></right>""") | |
| with gr.Row(): | |
| with gr.Column(): | |
| youtube_url = gr.Textbox(label="YouTube URL",value=f"{URL}") | |
| load_button = gr.Button("Load Video") | |
| output_win = gr.Video() | |
| with gr.Column(): | |
| with gr.Row(): | |
| cur_frame = gr.Number(label="Current Frame") | |
| fps_frames = gr.Number(label="Video FPS",interactive=False) | |
| total_frames = gr.Number(label="Total Frames",interactive=False) | |
| #run_button = gr.Button() | |
| with gr.Row(): | |
| bk = gr.Button("<") | |
| pl = gr.Button("Play") | |
| st = gr.Button("Stop") | |
| fw = gr.Button(">") | |
| det_win = gr.Image(source="webcam", streaming=True) | |
| with gr.Row(): | |
| pl_tog=gr.Number(visible=False) | |
| ins_cnt=gr.Number(visible=False) | |
| pl.click(tog_on,None,[pl_tog,play_state],show_progress=False) | |
| st.click(tog_off,None,[pl_tog,play_state],show_progress=False) | |
| pl_tog.change(pl_fn,[output_win,cur_frame,total_frames,fps_frames,pl_tog],[det_win,cur_frame],show_progress=False) | |
| cur_frame.change(pl_fn,[output_win,cur_frame,total_frames,fps_frames,pl_tog],[det_win,cur_frame],show_progress=False) | |
| bk.click(bk_fn,[cur_frame],cur_frame,show_progress=False) | |
| fw.click(fw_fn,[cur_frame,total_frames],cur_frame,show_progress=False) | |
| load_button.click(load,youtube_url,[output_win,cur_frame,total_frames,fps_frames,pl_tog]) | |
| #run_button.click(vid_play, [output_win,cur_frame], det_win) | |
| app.queue(concurrency_count=10).launch() |