Optical Flow Tracking光流追踪

输入：图片序列文件夹地址
输出：optical flow光流视频

import cv2
import numpy as np
import os
from os.path import isfile, join
import time

def getFlow(address1, address2,size):

    img1 = cv2.imread(address1)
    img2 = cv2.imread(address2)
    img1 = cv2.resize(img1, size)
    img2 = cv2.resize(img2, size)
    hsv = np.zeros_like(img1)
    hsv[...,1] = 255

    img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

    flow = cv2.calcOpticalFlowFarneback(
            img1,img2,
            pyr_scale=0.5,
            levels=7,
            winsize=5,
            iterations=5,
            poly_n=5,
            poly_sigma=11,
            flow=None,
            flags=0)

    mag, ang = cv2.cartToPolar(flow[...,0], flow[...,1])
    hsv[...,0] = ang*180/np.pi/2
    hsv[...,2] = cv2.normalize(mag,None,0,255,cv2.NORM_MINMAX)
    rgb = cv2.cvtColor(hsv,cv2.COLOR_HSV2BGR)
    return rgb

if __name__ == "__main__":
    #input and output address
    pathIn= './HighFrameRateTest2Sample100/Images/'
    pathOut = 'flow_tracking.avi'

    # preset output fps
    fps = 5

    # allocate list of images with ORB features
    frame_array = []
    files = [f for f in os.listdir(pathIn) if isfile(join(pathIn, f))]
    #for sorting the file names properly
    # files.sort(key = lambda x: x[5:-4])
    files.sort()
    size = (1024,512)
    for i in range(len(files) - 1):
        if files[i] != '.DS_Store':
            start = time.time()
            filename1=pathIn + files[i]
            filename2=pathIn + files[i+1]
            flow = getFlow(filename1,filename2,size)
            frame_array.append(flow)
            end = time.time()
            print('processing the ' + str(i) + ' image, processing time: ' + str(end - start))

    out = cv2.VideoWriter(pathOut,cv2.VideoWriter_fourcc(*'DIVX'), fps, size)
    for i in range(len(frame_array)):
        # writing to a image array
        out.write(frame_array[i])
    out.release()