video_avi_flow.py 文件源码

def extract_optical_flow(fn, times, frames=8, scale_factor=1.0):
    cap = cv2.VideoCapture(fn)
    if not cap.isOpened():
        return []
    n_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
    outputs = []
    if n_frames < frames * 2:
        return outputs

    def resize(im):
        if scale_factor != 1.0:
            new_size = (int(im.shape[1] * scale_factor), int(im.shape[0] * scale_factor))
            return cv2.resize(im, new_size, interpolation=cv2.INTER_LINEAR)
        else:
            return im

    for t in times:
        cap.set(cv2.CAP_PROP_POS_FRAMES, min(t * n_frames, n_frames - 1 - frames))
        ret, frame0 = cap.read()
        im0 = resize(cv2.cvtColor(frame0, cv2.COLOR_BGR2GRAY))
        mags = []
        middle_frame = frame0
        flows = []
        for f in range(frames - 1):
            ret, frame1 = cap.read()
            if f == frames // 2:
                middle_frame = frame1
            im1 = resize(cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY))
            flow = cv2.calcOpticalFlowFarneback(im0, im1,
                        None,
                        0.5, # py_scale
                        8,   # levels
                        int(40 * scale_factor),  # winsize
                        10,   # iterations
                        5,  # poly_n
                        1.1, # poly_sigma
                        cv2.OPTFLOW_FARNEBACK_GAUSSIAN)
            #mag, ang = cv2.cartToPolar(flow[...,0], flow[...,1])
            #mags.append(mag)
            flows.append(flow)
            im0 = im1
        flow = (np.mean(flows, 0) / 100).clip(-1, 1)

        #flow = np.mean(flows, 0)
        #flow /= (flow.mean() * 5 + 1e-5)
        #flow = flow.clip(-1, 1)
        #flows = flows / (np.mean(flows, 0, keepdims=True) + 1e-5)
        x = middle_frame[..., ::-1].astype(np.float32) / 255
        outputs.append((x, flow))
    return outputs