app and requriements

app.py

@@ -0,0 +1,107 @@
+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()

requirements.txt

@@ -0,0 +1,5 @@
+opencv-python
+onnx
+onnxruntime
+onnxruntime-gpu
+imread-from-url