|
|
import argparse |
|
|
import cv2 |
|
|
|
|
|
enable_cv2 = True |
|
|
|
|
|
import numpy as np |
|
|
import matplotlib.pyplot as plt |
|
|
from axengine import InferenceSession |
|
|
|
|
|
def bilinear_resize_numpy(array, new_h, new_w): |
|
|
h, w = array.shape |
|
|
x_ratio = w / new_w |
|
|
y_ratio = h / new_h |
|
|
|
|
|
resized = np.zeros((new_h, new_w), dtype=array.dtype) |
|
|
|
|
|
for i in range(new_h): |
|
|
for j in range(new_w): |
|
|
x = j * x_ratio |
|
|
y = i * y_ratio |
|
|
|
|
|
x_floor = int(x) |
|
|
y_floor = int(y) |
|
|
x_ceil = min(x_floor + 1, w - 1) |
|
|
y_ceil = min(y_floor + 1, h - 1) |
|
|
|
|
|
dx = x - x_floor |
|
|
dy = y - y_floor |
|
|
|
|
|
a = array[y_floor, x_floor] |
|
|
b = array[y_floor, x_ceil] |
|
|
c = array[y_ceil, x_floor] |
|
|
d = array[y_ceil, x_ceil] |
|
|
|
|
|
resized[i, j] = a * (1 - dx) * (1 - dy) + b * dx * (1 - dy) + c * (1 - dx) * dy + d * dx * dy |
|
|
|
|
|
return resized |
|
|
|
|
|
|
|
|
def resize_disp(disp, target_width, target_height, use_cv2=True): |
|
|
if use_cv2: |
|
|
disp = cv2.resize(disp, (target_width, target_height)) |
|
|
else: |
|
|
|
|
|
disp = bilinear_resize_numpy(disp, target_height, target_width) |
|
|
|
|
|
return disp |
|
|
|
|
|
def load_and_preprocess_image(image_path, target_width, target_height, use_cv2=True): |
|
|
|
|
|
if use_cv2: |
|
|
img = cv2.imread(image_path) |
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) |
|
|
orig_height, orig_width = img.shape[:2] |
|
|
img_resized = cv2.resize(img, (target_width,target_height) ) |
|
|
img_batch = img_resized[None] |
|
|
else: |
|
|
img = Image.open(image_path).convert('RGB') |
|
|
orig_width, orig_height = img.size |
|
|
img_resized = img.resize((target_width, target_height)) |
|
|
img_array = np.array(img_resized) |
|
|
img_batch = img_array[None] |
|
|
|
|
|
return img_batch, (orig_height, orig_width) |
|
|
|
|
|
def parse_args() -> argparse.Namespace: |
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument( |
|
|
"--left", |
|
|
type=str, |
|
|
required=True, |
|
|
help="Path to left image.", |
|
|
) |
|
|
parser.add_argument( |
|
|
"--right", |
|
|
type=str, |
|
|
required=True, |
|
|
help="Path to right image.", |
|
|
) |
|
|
parser.add_argument( |
|
|
"--model", |
|
|
type=str, |
|
|
required=True, |
|
|
help="Path to ONNX model.", |
|
|
) |
|
|
parser.add_argument( |
|
|
"--width", |
|
|
type=int, |
|
|
required=True, |
|
|
help="Width of input image.", |
|
|
) |
|
|
parser.add_argument( |
|
|
"--height", |
|
|
type=int, |
|
|
required=True, |
|
|
help="Height of input image.", |
|
|
) |
|
|
|
|
|
return parser.parse_args() |
|
|
|
|
|
|
|
|
def infer(left: str, right: str, model: str, width:int, height:int): |
|
|
|
|
|
image_left, (orig_h_left, orig_w_left) = load_and_preprocess_image(left, width, height, use_cv2=enable_cv2) |
|
|
image_right, (orig_h_right, orig_w_right) = load_and_preprocess_image(right, width, height, use_cv2=enable_cv2) |
|
|
|
|
|
assert orig_h_left == orig_h_right and orig_w_left == orig_w_right |
|
|
|
|
|
session = InferenceSession.load_from_model(model) |
|
|
flow_up = session.run(input_feed={"x1":image_left, "x2":image_right})["output"] |
|
|
|
|
|
|
|
|
flow_up = resize_disp(flow_up[0,0], orig_w_left, orig_h_left, use_cv2=enable_cv2) |
|
|
|
|
|
flow_up *= orig_w_left/width |
|
|
output = np.abs(flow_up) |
|
|
|
|
|
plt.imsave(f"output-ax.png", output, cmap='jet') |
|
|
|
|
|
return output |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
args = parse_args() |
|
|
infer(**vars(args)) |
|
|
|
