app.py

import cv2
import time
import numpy as np
import onnx
import onnxruntime
import os 
os.system('pip install --upgrade --force-reinstall onnxruntime')

# Ref: https://github.com/liruoteng/OpticalFlowToolkit/blob/5cf87b947a0032f58c922bbc22c0afb30b90c418/lib/flowlib.py#L249

import numpy as np

UNKNOWN_FLOW_THRESH = 1e7

def make_color_wheel():
    """
    Generate color wheel according Middlebury color code
    :return: Color wheel
    """
    RY = 15
    YG = 6
    GC = 4
    CB = 11
    BM = 13
    MR = 6

    ncols = RY + YG + GC + CB + BM + MR

    colorwheel = np.zeros([ncols, 3])

    col = 0

    # RY
    colorwheel[0:RY, 0] = 255
    colorwheel[0:RY, 1] = np.transpose(np.floor(255*np.arange(0, RY) / RY))
    col += RY

    # YG
    colorwheel[col:col+YG, 0] = 255 - np.transpose(np.floor(255*np.arange(0, YG) / YG))
    colorwheel[col:col+YG, 1] = 255
    col += YG

    # GC
    colorwheel[col:col+GC, 1] = 255
    colorwheel[col:col+GC, 2] = np.transpose(np.floor(255*np.arange(0, GC) / GC))
    col += GC

    # CB
    colorwheel[col:col+CB, 1] = 255 - np.transpose(np.floor(255*np.arange(0, CB) / CB))
    colorwheel[col:col+CB, 2] = 255
    col += CB

    # BM
    colorwheel[col:col+BM, 2] = 255
    colorwheel[col:col+BM, 0] = np.transpose(np.floor(255*np.arange(0, BM) / BM))
    col += + BM

    # MR
    colorwheel[col:col+MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR))
    colorwheel[col:col+MR, 0] = 255

    return colorwheel

colorwheel = make_color_wheel()

def compute_color(u, v):
    """
    compute optical flow color map
    :param u: optical flow horizontal map
    :param v: optical flow vertical map
    :return: optical flow in color code
    """
    [h, w] = u.shape
    img = np.zeros([h, w, 3])
    nanIdx = np.isnan(u) | np.isnan(v)
    u[nanIdx] = 0
    v[nanIdx] = 0

    ncols = np.size(colorwheel, 0)

    rad = np.sqrt(u**2+v**2)

    a = np.arctan2(-v, -u) / np.pi

    fk = (a+1) / 2 * (ncols - 1) + 1

    k0 = np.floor(fk).astype(int)

    k1 = k0 + 1
    k1[k1 == ncols+1] = 1
    f = fk - k0

    for i in range(0, np.size(colorwheel,1)):
        tmp = colorwheel[:, i]
        col0 = tmp[k0-1] / 255
        col1 = tmp[k1-1] / 255
        col = (1-f) * col0 + f * col1

        idx = rad <= 1
        col[idx] = 1-rad[idx]*(1-col[idx])
        notidx = np.logical_not(idx)

        col[notidx] *= 0.75
        img[:, :, i] = np.uint8(np.floor(255 * col*(1-nanIdx)))

    return img

def flow_to_image(flow):
    """
    Convert flow into middlebury color code image
    :param flow: optical flow map
    :return: optical flow image in middlebury color
    """
    u = flow[:, :, 0]
    v = flow[:, :, 1]

    maxu = -999.
    maxv = -999.
    minu = 999.
    minv = 999.

    idxUnknow = (abs(u) > UNKNOWN_FLOW_THRESH) | (abs(v) > UNKNOWN_FLOW_THRESH)
    u[idxUnknow] = 0
    v[idxUnknow] = 0

    maxu = max(maxu, np.max(u))
    minu = min(minu, np.min(u))

    maxv = max(maxv, np.max(v))
    minv = min(minv, np.min(v))

    rad = np.sqrt(u ** 2 + v ** 2)
    maxrad = max(-1, np.max(rad))

    u = u/(maxrad + np.finfo(float).eps)
    v = v/(maxrad + np.finfo(float).eps)

    img = compute_color(u, v)

    idx = np.repeat(idxUnknow[:, :, np.newaxis], 3, axis=2)
    img[idx] = 0

    return np.uint8(img)
class Raft():

	def __init__(self, model_path):

		# Initialize model
		self.initialize_model(model_path)

	def __call__(self, img1, img2):

		return self.estimate_flow(img1, img2)

	def initialize_model(self, model_path):

		self.session = onnxruntime.InferenceSession(model_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

		# Get model info
		self.get_input_details()
		self.get_output_details()

	def estimate_flow(self, img1, img2):

		input_tensor1 = self.prepare_input(img1)
		input_tensor2 = self.prepare_input(img2)

		outputs = self.inference(input_tensor1, input_tensor2)
		
		self.flow_map = self.process_output(outputs)

		return self.flow_map

	def prepare_input(self, img):

		img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

		self.img_height, self.img_width = img.shape[:2]

		img_input = cv2.resize(img, (self.input_width,self.input_height))

		# img_input = img_input/255
		img_input = img_input.transpose(2, 0, 1)
		img_input = img_input[np.newaxis,:,:,:]        

		return img_input.astype(np.float32)

	def inference(self, input_tensor1, input_tensor2):

		# start = time.time()
		outputs = self.session.run(self.output_names, {self.input_names[0]: input_tensor1, 
													   self.input_names[1]: input_tensor2})

		# print(time.time() - start)
		return outputs

	def process_output(self, output): 

		flow_map = output[1][0].transpose(1, 2, 0)

		return flow_map

	def draw_flow(self):

		# Convert flow to image
		flow_img = flow_to_image(self.flow_map)

		# Convert to BGR
		flow_img = cv2.cvtColor(flow_img, cv2.COLOR_RGB2BGR)

		# Resize the depth map to match the input image shape
		return cv2.resize(flow_img, (self.img_width,self.img_height))

	def get_input_details(self):

		model_inputs = self.session.get_inputs()
		self.input_names = [model_inputs[i].name for i in range(len(model_inputs))]

		self.input_shape = model_inputs[0].shape
		self.input_height = self.input_shape[2]
		self.input_width = self.input_shape[3]

	def get_output_details(self):

		model_outputs = self.session.get_outputs()
		self.output_names = [model_outputs[i].name for i in range(len(model_outputs))]

		self.output_shape = model_outputs[0].shape
		self.output_height = self.output_shape[2]
		self.output_width = self.output_shape[3]

if __name__ == '__main__':
	
	from imread_from_url import imread_from_url

	# Initialize model
	model_path='raft_small_iter10_240x320.onnx'
	flow_estimator = Raft(model_path)

	# Read inference image
	img1 = imread_from_url("https://github.com/princeton-vl/RAFT/blob/master/demo-frames/frame_0016.png?raw=true")
	img2 = imread_from_url("https://github.com/princeton-vl/RAFT/blob/master/demo-frames/frame_0025.png?raw=true")

	# Estimate flow and colorize it
	flow_map = flow_estimator(img1, img2)
	flow_img = flow_estimator.draw_flow()

	combined_img = np.hstack((img1, img2, flow_img))

	#cv2.namedWindow("Estimated flow", cv2.WINDOW_NORMAL)
	#cv2.imshow("Estimated flow", combined_img)
	#cv2.waitKey(0)

import os
import cv2
import gradio as gr
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 = 'models/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()