epic_flow.py 文件源码

def compute_flow(impath1, impath2, outdir,
                    fbcodepath=os.getenv("HOME") + '/fbcode'):
    stem = os.path.splitext(os.path.basename(impath1))[0]
    deepmatch_cmd = os.path.join(fbcodepath,
                                    '_bin/experimental/deeplearning/dpathak' +
                                    '/video-processing/deepmatch/deepmatch')
    call([deepmatch_cmd, impath1, impath2, '-out',
                os.path.join(outdir, stem + '_sparse.txt'), '-downscale', '2'])
    img1 = cv2.imread(impath1).astype(float)
    M = np.zeros((img1.shape[0], img1.shape[1]), dtype=np.float32)
    filt = np.array([[1., -1.]]).reshape((1, -1))
    for c in range(3):
        gx = convolve2d(img1[:, :, c], filt, mode='same')
        gy = convolve2d(img1[:, :, c], filt.T, mode='same')
        M = M + gx**2 + gy**2

    M = M / np.max(M)
    with open(os.path.join(outdir, '_edges.bin'), 'w') as f:
        M.tofile(f)

    epicflow_command = os.path.join(fbcodepath,
                                    '_bin/experimental/deeplearning/dpathak' +
                                    '/video-processing/epicflow/epicflow')
    call([epicflow_command, impath1, impath2,
                os.path.join(outdir, '_edges.bin'),
                os.path.join(outdir, stem + '_sparse.txt'),
                os.path.join(outdir, 'flow.flo')])

    flow = read_flo(os.path.join(outdir, 'flow.flo'))
    hsv = np.zeros_like(img1).astype(np.uint8)
    hsv[..., 1] = 255
    mag, ang = cv2.cartToPolar(flow[..., 0].astype(float),
                                flow[..., 1].astype(float))
    hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
    hsv[..., 0] = ang * 180 / np.pi / 2
    bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
    cv2.imwrite(os.path.join(outdir, stem + '_flow.png'), bgr)