Spaces:
Runtime error
Runtime error
File size: 3,235 Bytes
af5763a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
import os
import cv2
import gradio as gr
from raft import Raft
import yt_dlp
def download_youtube_video(youtube_url, output_filename):
ydl_opts = {
'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best',
'outtmpl': output_filename,
}
with yt_dlp.YoutubeDL(ydl_opts) as ydl:
ydl.download([youtube_url])
def process_video(youtube_url, start_time, flow_frame_offset):
model_path = 'raft_small_iter10_240x320.onnx'
flow_estimator = Raft(model_path)
output_filename = 'downloaded_video.mp4'
processed_output = 'processed_video.mp4'
# Download video
if os.path.exists(output_filename):
os.remove(output_filename)
download_youtube_video(youtube_url, output_filename)
cap = cv2.VideoCapture(output_filename)
if not cap.isOpened():
return "Error: Could not open video."
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter(processed_output, fourcc, fps, (frame_width, frame_height))
cap.set(cv2.CAP_PROP_POS_FRAMES, start_time * fps)
frame_list = []
frame_num = 0
while cap.isOpened():
ret, prev_frame = cap.read()
if not ret:
break
frame_list.append(prev_frame)
frame_num += 1
if frame_num <= flow_frame_offset:
continue
flow_map = flow_estimator(frame_list[0], frame_list[-1])
flow_img = flow_estimator.draw_flow()
alpha = 0.5
combined_img = cv2.addWeighted(frame_list[0], alpha, flow_img, (1 - alpha), 0)
if combined_img is None:
break
out.write(combined_img)
frame_list.pop(0)
cap.release()
out.release()
return processed_output
examples = [
["https://www.youtube.com/watch?v=is38pqgbj6A", 5, 50, "output_1.mp4"],
["https://www.youtube.com/watch?v=AdbrfoxiAtk", 0, 60, "output_2.mp4"],
["https://www.youtube.com/watch?v=vWGg0iPmI8k", 13, 70, "output_3.mp4"],
]
with gr.Blocks() as app:
gr.Markdown("### Optical Flow Video Processing\n"
"Enter a YouTube URL, set the start time and flow frame offset, "
"then click 'Process Video' to see the optical flow processing.")
with gr.Row():
with gr.Column():
youtube_url = gr.Textbox(label="YouTube URL", placeholder="Enter YouTube Video URL Here")
start_time = gr.Slider(minimum=0, maximum=60, label="Start Time (seconds)", step=1)
flow_frame_offset = gr.Slider(minimum=1, maximum=100, label="Flow Frame Offset", step=1)
submit_button = gr.Button("Process Video")
with gr.Column():
output_video = gr.Video(label="Processed Video")
submit_button.click(
fn=process_video,
inputs=[youtube_url, start_time, flow_frame_offset],
outputs=output_video
)
gr.Examples(examples=examples,
inputs=[youtube_url, start_time, flow_frame_offset],
fn=process_video,
outputs=output_video,
cache_examples=False)
app.launch()
|