OpenCV Optical Flow

gistfile1.txt

import glob
images = glob.glob("*.jpg")
import re
import os
import matplotlib.pyplot as plt

def natural_sort(l): 
    convert = lambda text: int(text) if text.isdigit() else text.lower() 
    alphanum_key = lambda key: [ convert(c) for c in re.split('([0-9]+)', key) ] 
    return sorted(l, key = alphanum_key)

images = natural_sort(images)

import numpy as np
import cv2
from PIL import Image


MAX_CORNERS = 850

feature_params = dict( maxCorners = MAX_CORNERS,
                       qualityLevel = 0.005,
                       minDistance = 4,
                       blockSize = 7 )

lk_params = dict( winSize  = (15,15),
                  maxLevel = 2,
                  criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))


color = np.random.randint(0,255,(MAX_CORNERS,3))

old_gray = cv2.cvtColor(np.array(Image.open(images[0]))[55:,:,:], cv2.COLOR_BGR2GRAY)
mask = np.zeros_like(np.array(Image.open(images[0]))[55:,:,:])
p0 = cv2.goodFeaturesToTrack(old_gray, mask = None, **feature_params)

good_new = []
good_old = []

hard_reset_count = 0

for img_name in images:
    frame = np.array(Image.open(img_name))[55:,:,:]
    frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    # calculate optical flow
    p1, st, err = cv2.calcOpticalFlowPyrLK(old_gray, frame_gray, p0, None, **lk_params)

    #print(len(good_new),len(good_old))

    if p1 is not None:
        # Select good points
        good_new = p1[st==1]
        good_old = p0[st==1]

        # draw the tracks
        lift_count_active = 0
        for i,(new,old) in enumerate(zip(good_new,good_old)):
            a,b = new.ravel()
            c,d = old.ravel()
            dist = np.linalg.norm(np.array(a,b)-np.array(c,d))
            #if d - b > 5 and dist > 2 and dist < 15:
            if dist > 2.0 and dist < 7:
                lift_count_active += 1
                mask = cv2.line(mask, (a,b),(c,d), color[i].tolist(), 2)
                frame = cv2.circle(frame,(a,b),5,color[i].tolist(),-1)

        if lift_count_active > 40:
            print("lift_count_active :",lift_count_active)
        img = cv2.add(frame,mask)

            # Now update the previous frame and previous points
        old_gray = frame_gray.copy()
        p0 = good_new.reshape(-1,1,2)
    else :
        img = frame
        good_new = []
        good_old = []
        old_gray = frame_gray.copy()
        mask = np.zeros_like(np.array(Image.open(images[0]))[55:,:,:])
        p0 = cv2.goodFeaturesToTrack(old_gray, mask = None, **feature_params)

    if hard_reset_count > 15:
        hard_reset_count = 0
        img = frame
        good_new = []
        good_old = []
        old_gray = frame_gray.copy()
        mask = np.zeros_like(np.array(Image.open(images[0]))[55:,:,:])
        p0 = cv2.goodFeaturesToTrack(old_gray, mask = None, **feature_params)


    cv2.imshow('frame',img)
    
    #while True:
    #    k = cv2.waitKey(30) & 0xff
    #    if k == 32:
    #        break  
    fig = plt.figure(figsize=(18,9))
    ax1 = fig.add_subplot(1, 2, 1)
    ax2 = fig.add_subplot(1, 2, 2)
    ax1.imshow(np.array(Image.open(img_name))[55:,:,:])
    ax2.imshow(img)
    plt.savefig(img_name+"_optical_flow.png", dpi=150, transparent=True, bbox_inches='tight')
    plt.close()

    if k == 27:
        break

    hard_reset_count += 1