Project 8

import cv2
import numpy as np
from matplotlib import pyplot as plt

def display_img(disp_img):
    cv2.imshow('img_display',disp_img)
    cv2.waitKey(0)

def read_img(path):
    img_matrix = cv2.imread(path)
    display_img(img_matrix)
    return img_matrix

def img_cvt2gray(img):
    gray_img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    return gray_img

def lk_opticalFlow(prev_img, curr_img, cornerPts):
    Lk_pts, status, err = cv2.calcOpticalFlowPyrLK(prev_img, curr_img, cornerPts, None,      
             winSize  = (15, 15), maxLevel = 2, 
             criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))
    return(Lk_pts, status, err)

def cornerPoints(grey):
    # params for ShiTomasi corner detection
    points = cv2.goodFeaturesToTrack(grey, mask = None ,maxCorners = 100, 
             qualityLevel = 0.3, minDistance = 7, blockSize = 7)
    return points

cap = cv2.VideoCapture("assignment_8/traffic_assignment.mp4")

# Write video dimension
frame_width = int(cap.get(3))
frame_height = int(cap.get(4)) 
size = (frame_width, frame_height)
# save video in .avi format at 30fps
result_opticalFlow = cv2.VideoWriter('assignment_8/Traffic_opticalFlow.avi', cv2.VideoWriter_fourcc(*'MJPG'), 30, size)
# save video in .avi format at 30fps
result_mask = cv2.VideoWriter('assignment_8/Traffic_mask.avi', cv2.VideoWriter_fourcc(*'MJPG'), 30, size)
# Create some random colors
color = np.random.randint(0, 255, (100, 3)) #100 array of dim ((1, 3 )- (R,G,B)) with random val between 0 and 255
# Take first frame and find corners in it
ret, old_frame = cap.read()
old_gray = img_cvt2gray(old_frame)
Prev_points = cornerPoints(old_gray) # detect initial location of corner points using ShiTomasi corner detection algorithm
# Create a mask image for drawing purposes
mask = np.zeros_like(old_frame)
once=True
while(1):
    ret, frame = cap.read()
    if not ret:
        print('No frames grabbed!')
        break
    frame_gray = img_cvt2gray(frame)
    
    # Track detected corners using Lucas kanade sparse optical flow and returns next    displacement of corner points using current position
    (Current_points,status_Lk,Error_Lk) = lk_opticalFlow(old_gray,frame_gray,Prev_points)
    # Select good points
    if Current_points is not None:
        Lk_good_new = Current_points[status_Lk==1] # converts 3d array to 2d
        Lk_good_old = Prev_points[status_Lk==1]       
    # draw the tracks
    for i, (new, old) in enumerate(zip(Lk_good_new, Lk_good_old)):
        a, b = new.ravel()
        c, d = old.ravel()
        mask = cv2.line(mask, (int(a), int(b)), (int(c), int(d)), color[i].tolist(), 2)
        frame = cv2.circle(frame, (int(a), int(b)), 5, color[i].tolist(), -1)
    
    #merge mask and frmae
    img_opticalFlow = cv2.add(frame, mask)
    #save video
    result_opticalFlow.write(img_opticalFlow)
    result_mask.write(mask)
    #display video frame
    cv2.imshow('optical flow mask',mask)
    cv2.imshow('optical flow',img_opticalFlow)
    cv2.waitKey(1)
    # Now update the previous frame and previous points
    old_gray = frame_gray.copy()
    Prev_points = Lk_good_new.reshape(-1, 1, 2) # reshape to (1,2) dim matrix of n matrix (-1 means, max matrix possible)
cv2.destroyAllWindows()