def infer()

in optical_flow_flownet2_homography.py [0:0]

• 43 lines of code
• 2 McCabe index (conditional complexity)

def infer(args, Flownet, device, img1_name, img2_name):
    img1, img2, img2_reg, H_BA = getimage(img1_name, img2_name)
    _, imgH, imgW = img1.shape
    img1 = img1[None, :, :]
    img2 = img2[None, :, :]
    img2_reg = img2_reg[None, :, :]
    img1 = img1.to(device)
    img2 = img2.to(device)
    img2_reg = img2_reg.to(device)

    if args.homography != 1:
        sz = img1.size()
        img1_view = img1.view(sz[0], sz[1], 1, sz[2], sz[3])
        img2_view = img2.view(sz[0], sz[1], 1, sz[2], sz[3])
        inputs = torch.cat((img1_view, img2_view), dim=2)
        flow = Flownet(inputs)[0].permute(1, 2, 0).data.cpu().numpy()
    else:
        sz = img1.size()
        img1_view = img1.view(sz[0], sz[1], 1, sz[2], sz[3])
        img2_reg_view = img2_reg.view(sz[0], sz[1], 1, sz[2], sz[3])
        inputs = torch.cat((img1_view, img2_reg_view), dim=2)
        flow = Flownet(inputs)[0].permute(1, 2, 0).data.cpu().numpy()

        (fy, fx) = np.mgrid[0:imgH, 0:imgW].astype(np.float32)

        fxx = copy.deepcopy(fx) + flow[:, :, 0]
        fyy = copy.deepcopy(fy) + flow[:, :, 1]

        (fxxx, fyyy, fz) = np.linalg.inv(H_BA).dot(
            np.concatenate(
                (
                    fxx.reshape(1, -1),
                    fyy.reshape(1, -1),
                    np.ones_like(fyy).reshape(1, -1),
                ),
                axis=0,
            )
        )
        fxxx, fyyy = fxxx / fz, fyyy / fz

        flow = np.concatenate(
            (
                fxxx.reshape(imgH, imgW, 1) - fx.reshape(imgH, imgW, 1),
                fyyy.reshape(imgH, imgW, 1) - fy.reshape(imgH, imgW, 1),
            ),
            axis=2,
        )

    return flow