光流估计

# -*- coding:utf-8 -*-
import cv2
import numpy as np

def video_operate():
    capture = cv2.VideoCapture("test.avi")
    feature_param  = dict(maxCorners =100,qualityLevel = 0.3,minDistance= 7)
    lk_param = dict(winSize=(15,15),maxLevel=2)
    color_ = np.random.randint(0,255,(100,3))
    ret,old_frame = capture.read()
    old_gray =cv2.cvtColor(old_frame,cv2.COLOR_BGR2GRAY)
    p0 = cv2.goodFeaturesToTrack(old_gray,mask=None,**feature_param)
    mask = np.zeros_like(old_frame)

    while True:
        ret,frame = capture.read()
        frame_gray = cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)
        p1,st,err = cv2.calcOpticalFlowPyrLK(old_gray,frame_gray,p0,None,**lk_param)
        good_new = p1[st==1]
        good_old = p0[st==1]
        for i,(new,old) in enumerate(zip(good_new,good_old)):
            a,b = new.ravel()
            c,d = old.ravel()
            mask = cv2.line(mask,(a,b),(c,d),color_[i].tolist(),2)
            frame = cv2.circle(frame,(a,b),5,color_[i].tolist(),-1)

        img = cv2.add(frame,mask)
        cv2.imshow('frame',img)
        key = cv2.waitKey(150) & 0xFF
        if key == 27:
            break
        old_gray = frame_gray.copy()
        p0= good_new.reshape(-1,1,2)

    capture.release()
    cv2.destroyAllWindows()
def main():
    video_operate()
if __name__  == "__main__":
    main()