Ejemplo n.º 1
0
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
Ejemplo n.º 2
0
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)
Ejemplo n.º 3
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)