optflow_utils.py 文件源码

def dense_optical_flow(im1, im2, pyr_scale=0.5, levels=3, winsize=5, 
                       iterations=3, poly_n=5, poly_sigma=1.2, fb_threshold=-1, 
                       mask1=None, mask2=None, 
                       flow1=None, flow2=None): 

    if flow1 is None: 
        fflow = cv2.calcOpticalFlowFarneback(to_gray(im1), to_gray(im2), pyr_scale, levels, winsize, 
                                             iterations, poly_n, poly_sigma, 0)
    else: 
        fflow = cv2.calcOpticalFlowFarneback(to_gray(im1), to_gray(im2), pyr_scale, levels, winsize, 
                                             iterations, poly_n, poly_sigma, 0, flow1.copy())

    if mask1 is not None: 
        fflow[~mask1.astype(np.bool)] = np.nan

    if fb_threshold > 0: 
        H, W = im1.shape[:2]
        xs, ys = np.meshgrid(np.arange(W), np.arange(H))
        xys1 = np.dstack([xs, ys])
        xys2 = xys1 + fflow
        rflow = dense_optical_flow(im2, im1, pyr_scale=pyr_scale, levels=levels, 
                                   winsize=winsize, iterations=iterations, poly_n=poly_n, 
                                   poly_sigma=poly_sigma, fb_threshold=-1)
        if mask2 is not None: 
            rflow[~mask2.astype(np.bool)] = np.nan

        xys1r = xys2 + rflow
        fb_bad = (np.fabs(xys1r - xys1) > fb_threshold).all(axis=2)
        fflow[fb_bad] = np.nan

    return fflow