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 get_data(self):
for flow_path in self.flows:
input_path = flow_path.replace(
self.path_prefix, os.path.join(
self.data_path,
'clean',
))
frame_id = int(input_path[-8:-4])
input_a_path = '%s%04i.png' % (input_path[:-8], frame_id)
input_b_path = '%s%04i.png' % (input_path[:-8], frame_id + 1)
input_a = cv2.imread(input_a_path)
input_b = cv2.imread(input_b_path)
flow = Flow.read(flow_path)
# most implementation just crop the center
# which seems to be accepted practise
h, w = input_a.shape[:2]
h_ = (h // 64) * 64
w_ = (w // 64) * 64
h_start = (h - h_) // 2
w_start = (w - w_) // 2
# this is ugly
h_end = -h_start if h_start > 0 else h
w_end = -w_start if w_start > 0 else w
input_a = input_a[h_start:h_end, w_start:w_end, :]
input_b = input_b[h_start:h_end, w_start:w_end, :]
flow = flow[h_start:h_end, w_start:w_end, :]
yield [input_a, input_b, flow]
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)
def __iter__(self):
for flow_path in self.flows:
input_path = flow_path.replace(
self.path_prefix, os.path.join(self.data_path, 'clean'))
frame_id = int(input_path[-8:-4])
input_a_path = '%s%04i.png' % (input_path[:-8], frame_id)
input_b_path = '%s%04i.png' % (input_path[:-8], frame_id + 1)
input_a = cv2.imread(input_a_path)
input_b = cv2.imread(input_b_path)
flow = Flow.read(flow_path)
# most implementation just crop the center
# which seems to be accepted practise
h, w = input_a.shape[:2]
newh = (h // 64) * 64
neww = (w // 64) * 64
aug = imgaug.CenterCrop((newh, neww))
input_a = aug.augment(input_a)
input_b = aug.augment(input_b)
flow = aug.augment(flow)
yield [input_a, input_b, flow]