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)