def apply(model, model_path, images, ground_truth=None):
left = cv2.imread(images[0])
h, w = left.shape[:2]
newh = (h // 64) * 64
neww = (w // 64) * 64
aug = imgaug.CenterCrop((newh, neww))
left = aug.augment(left)
predict_func = OfflinePredictor(
PredictConfig(model=model(height=newh, width=neww),
session_init=get_model_loader(model_path),
input_names=['left', 'right'],
output_names=['prediction']))
for right in images[1:]:
right = aug.augment(cv2.imread(right))
left_input, right_input = [
x.astype('float32').transpose(2, 0, 1)[None, ...]
for x in [left, right]
]
output = predict_func(left_input, right_input)[0].transpose(0, 2, 3, 1)
flow = Flow()
img = flow.visualize(output[0])
patches = [left, right, img * 255.]
if ground_truth is not None:
patches.append(flow.visualize(Flow.read(ground_truth)) * 255.)
img = viz.stack_patches(patches, 2, 2)
cv2.imshow('flow output', img)
cv2.imwrite('flow_prediction.png', img)
cv2.waitKey(0)
left = right
def apply(model_path, left, right):
left = cv2.imread(left).astype(np.float32)
right = cv2.imread(right).astype(np.float32)
assert left.shape == right.shape
h_in, w_in = left.shape[:2]
# images needs to be divisible by 64
h = int(np.ceil(h_in / 64.) * 64.)
w = int(np.ceil(w_in / 64.) * 64.)
print('resize inputs (%i, %i) to (%i, %i)' % (h_in, w_in, h, w))
left = cv2.resize(left, (w, h)).transpose(2, 0, 1)[None, ...]
right = cv2.resize(right, (w, h)).transpose(2, 0, 1)[None, ...]
predict_func = OfflinePredictor(PredictConfig(
model=PWCModel(),
session_init=get_model_loader(model_path),
input_names=['left', 'right'],
output_names=['prediction']))
output = predict_func(left, right)[0].transpose(0, 2, 3, 1)[0]
dx = cv2.resize(output[:, :, 0], (w_in, h_in)) * w_in / float(w)
dy = cv2.resize(output[:, :, 1], (w_in, h_in)) * h_in / float(h)
output = np.dstack((dx, dy))
flow = Flow()
img = flow.visualize(output)
cv2.imwrite('pwc_output.png', img * 255)
cv2.imshow('flow output', img)
cv2.waitKey(0)
def apply(model_name, model_path, left, right, ground_truth=None):
model = MODEL_MAP[model_name]
left = cv2.imread(left).astype(np.float32).transpose(2, 0, 1)[None, ...]
right = cv2.imread(right).astype(np.float32).transpose(2, 0, 1)[None, ...]
predict_func = OfflinePredictor(
PredictConfig(model=model(),
session_init=get_model_loader(model_path),
input_names=['left', 'right'],
output_names=['prediction']))
output = predict_func(left, right)[0].transpose(0, 2, 3, 1)
flow = Flow()
img = flow.visualize(output[0])
if ground_truth is not None:
img = np.concatenate(
[img, flow.visualize(Flow.read(ground_truth))], axis=1)
cv2.imshow('flow output', img)
cv2.waitKey(0)