app file

app.py

@@ -1,3 +1,254 @@
+import cv2
+import time
+import numpy as np
+import onnx
+import 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
@@ -13,7 +264,7 @@ def download_youtube_video(youtube_url, output_filename):
         ydl.download([youtube_url])
 
 def process_video(youtube_url, start_time, flow_frame_offset):
-    model_path = 'raft_small_iter10_240x320.onnx'
+    model_path = 'models/raft_small_iter10_240x320.onnx'
     flow_estimator = Raft(model_path)
 
     output_filename = 'downloaded_video.mp4'