1+ # Usage
2+ # python scripts/optical_flow_sparse_manual.py -p 'resources/car_race_01.mp4'
3+ import datetime
4+ import sys
5+ import numpy as np
6+ import cv2
7+ import argparse
8+
9+ # Parse the arguments
10+ ap = argparse .ArgumentParser ()
11+ ap .add_argument ("-p" , "--path" , help = "Path to video file" , default = "resources/car_race_02.mp4" )
12+ args = vars (ap .parse_args ())
13+
14+ cap = cv2 .VideoCapture (args ["path" ])
15+ if not cap .isOpened ():
16+ print ("[ERROR] cannot open video file" )
17+ sys .exit ()
18+
19+ # generate initial corners of detected object
20+ # set limit, minimum distance in pixels and quality of object corner to be tracked
21+ parameters_shitomasi = dict (maxCorners = 100 , qualityLevel = 0.3 , minDistance = 7 )
22+ # set min size of tracked object, e.g. 15x15px
23+ parameter_lucas_kanade = dict (winSize = (15 , 15 ), maxLevel = 2 , criteria = (cv2 .TERM_CRITERIA_EPS |
24+ cv2 .TERM_CRITERIA_COUNT , 10 , 0.03 ))
25+ # create random colours for visualization for all 100 max corners for RGB channels
26+ colours = np .random .randint (0 , 255 , (100 , 3 ))
27+
28+ # get first video frame
29+ ok , frame = cap .read ()
30+ if not ok :
31+ print ("[ERROR] cannot get frame from video" )
32+ sys .exit ()
33+ # convert to grayscale
34+ frame_gray_init = cv2 .cvtColor (frame , cv2 .COLOR_BGR2GRAY )
35+
36+ # Use optical flow to detect object corners / edges from initial frame
37+ edges = cv2 .goodFeaturesToTrack (frame_gray_init , mask = None , ** parameters_shitomasi )
38+ # [Debug] show amount of found edges
39+ # max value = maxCorners see above. Reduce qualityLevel to get more hits
40+ # print(len(edges))
41+
42+ # create a black canvas the size of the initial frame
43+ canvas = np .zeros_like (frame )
44+
45+ # Optional recording parameter
46+ frame_width = int (cap .get (cv2 .CAP_PROP_FRAME_WIDTH ))
47+ frame_height = int (cap .get (cv2 .CAP_PROP_FRAME_HEIGHT ))
48+ fps = int (cap .get (cv2 .CAP_PROP_FPS ))
49+ video_codec = cv2 .VideoWriter_fourcc ('m' , 'p' , '4' , 'v' )
50+ prefix = 'recordings/' + datetime .datetime .now ().strftime ("%y%m%d_%H%M%S" )
51+ basename = "object_track.mp4"
52+ video_output = cv2 .VideoWriter ("_" .join ([prefix , basename ]), video_codec , fps , (frame_width , frame_height ))
53+
54+ # loop through the remaining frames of the video
55+ # and apply algorithm to track selected objects
56+ while True :
57+ # get next frame
58+ ok , frame = cap .read ()
59+ if not ok :
60+ print ("[INFO] end of file reached" )
61+ break
62+ # prepare grayscale image
63+ frame_gray = cv2 .cvtColor (frame , cv2 .COLOR_BGR2GRAY )
64+ # update object corners by comparing with found edges in initial frame
65+ update_edges , status , errors = cv2 .calcOpticalFlowPyrLK (frame_gray_init , frame_gray , edges , None ,
66+ ** parameter_lucas_kanade )
67+ # only update edges if algorithm successfully tracked
68+ new_edges = update_edges [status == 1 ]
69+ # to calculate directional flow we need to compare with previous position
70+ old_edges = edges [status == 1 ]
71+
72+ for i , (new , old ) in enumerate (zip (new_edges , old_edges )):
73+ a , b = new .ravel ()
74+ c , d = old .ravel ()
75+
76+ # draw line between old and new corner point with random colour
77+ mask = cv2 .line (canvas , (int (a ), int (b )), (int (c ), int (d )), colours [i ].tolist (), 2 )
78+ # draw circle around new position
79+ frame = cv2 .circle (frame , (int (a ), int (b )), 5 , colours [i ].tolist (), - 1 )
80+
81+ result = cv2 .add (frame , mask )
82+ # optional recording result/mask
83+ # video_output.write(result)
84+ cv2 .imshow ('Optical Flow (sparse)' , result )
85+ if cv2 .waitKey (1 ) & 0xFF == ord ('q' ):
86+ break
87+ # overwrite initial frame with current before restarting the loop
88+ frame_gray_init = frame_gray .copy ()
89+ # update to new edges before restarting the loop
90+ edges = new_edges .reshape (- 1 , 1 , 2 )
91+
92+
93+ cv2 .destroyAllWindows ()
94+ cap .release ()
0 commit comments