opticalflow_TVL-1.py

import numpy as np
import cv2 as cv


# Most of the code come from https://docs.opencv.org/master/d4/dee/tutorial_optical_flow.html
# But, The algorithm is changed to TVL-1
cap = cv.VideoCapture(cv.samples.findFile("vtest.avi"))
ret, frame1 = cap.read()
prvs = cv.cvtColor(frame1, cv.COLOR_BGR2GRAY)
hsv = np.zeros_like(frame1)
hsv[..., 1] = 255
while(1):
    ret, frame2 = cap.read()
    next = cv.cvtColor(frame2, cv.COLOR_BGR2GRAY)

    # TVL-1
    # opencv_python==4.1.2.30
    # opencv-contrib-python==4.1.2.30
    optical_flow = cv.optflow.DualTVL1OpticalFlow_create()
    flow = optical_flow.calc(prvs, next, None)

    mag, ang = cv.cartToPolar(flow[..., 0], flow[..., 1])

    hsv[..., 0] = ang*180/np.pi/2
    hsv[..., 2] = cv.normalize(mag, None, 0, 255, cv.NORM_MINMAX)

    bgr = cv.cvtColor(hsv, cv.COLOR_HSV2BGR)

    cv.imshow('frame2', bgr)
    k = cv.waitKey(30) & 0xff
    if k == 27:
        break
    elif k == ord('s'):
        cv.imwrite('opticalfb.png', frame2)
        cv.imwrite('opticalhsv.png', bgr)

    prvs = next