Optical Flow with Python

1. 개요

opencv에서 제공하는 함수를 이용하여 optical flow를 계산한다.
기본 뼈대 프로그램으로 활용하기에 편리하다.

2. 코드

import cv2, numpy as np

def drawOpticalFlow(img, flow,indices,hf,wf):  
    for x, y in indices:  
        #optocal flow 이동거리  
        dx, dy = flow[y, x].astype(np.int)  
        #이동거리 선그리기  
        cv2.line(img, (x, y), (x + dx, y + dy), (0, 255, 0), 1, cv2.LINE_AA)  
        cv2.line(img, (wf-100,hf), (wf+100,hf), (255, 0, 0), 1, cv2.LINE_AA)  
        cv2.line(img, (wf,hf-100), (wf,hf+100), (255, 0, 0), 1 , cv2.LINE_AA)  
        # 기준위치 원그리기  
        cv2.circle(img, (x, y), 1, (0, 0, 255), -1)

def main():  
    prev = None  
    #cap = cv2.VideoCapture('tello\_capture\_video.mp4')  
    cap = cv2.VideoCapture(0)

    ret, frame = cap.read()
    h, w = frame.shape[:2]
    hf, wf = int(h/2), int(w/2)

    # 이미지 평면 위의 기준위치
    step = 32
    idx_y, idx_x = np.mgrid[step/2:h:step, step/2:w:step].astype(np.int)
    indices = np.stack((idx_x, idx_y), axis=-1).reshape(-1, 2)

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret: 
            break
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
       if prev is None:
            prev = gray
            continue
       else:
            flow = cv2.calcOpticalFlowFarneback(prev, gray, None, \
                   0.5, 3, 15, 3, 5, 1.1, cv2.OPTFLOW_FARNEBACK_GAUSSIAN)
            drawOpticalFlow(frame, flow,indices,hf,wf)
            prev = gray

        cv2.imshow('Farneback', frame)
        if cv2.waitKey(1) == 27:
            break
    cap.release()
    cv2.destroyAllWindows()
    pass

if __name__ == '__main__':  
    main()