Home | History | Annotate | Download | only in python2 
      1 #!/usr/bin/env python
      2 
      3 '''
      4 Lucas-Kanade tracker
      5 ====================
      6 
      7 Lucas-Kanade sparse optical flow demo. Uses goodFeaturesToTrack
      8 for track initialization and back-tracking for match verification
      9 between frames.
     10 
     11 Usage
     12 -----
     13 lk_track.py [<video_source>]
     14 
     15 
     16 Keys
     17 ----
     18 ESC - exit
     19 '''
     20 
     21 import numpy as np
     22 import cv2
     23 import video
     24 from common import anorm2, draw_str
     25 from time import clock
     26 
     27 lk_params = dict( winSize  = (15, 15),
     28                   maxLevel = 2,
     29                   criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))
     30 
     31 feature_params = dict( maxCorners = 500,
     32                        qualityLevel = 0.3,
     33                        minDistance = 7,
     34                        blockSize = 7 )
     35 
     36 class App:
     37     def __init__(self, video_src):
     38         self.track_len = 10
     39         self.detect_interval = 5
     40         self.tracks = []
     41         self.cam = video.create_capture(video_src)
     42         self.frame_idx = 0
     43 
     44     def run(self):
     45         while True:
     46             ret, frame = self.cam.read()
     47             frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
     48             vis = frame.copy()
     49 
     50             if len(self.tracks) > 0:
     51                 img0, img1 = self.prev_gray, frame_gray
     52                 p0 = np.float32([tr[-1] for tr in self.tracks]).reshape(-1, 1, 2)
     53                 p1, st, err = cv2.calcOpticalFlowPyrLK(img0, img1, p0, None, **lk_params)
     54                 p0r, st, err = cv2.calcOpticalFlowPyrLK(img1, img0, p1, None, **lk_params)
     55                 d = abs(p0-p0r).reshape(-1, 2).max(-1)
     56                 good = d < 1
     57                 new_tracks = []
     58                 for tr, (x, y), good_flag in zip(self.tracks, p1.reshape(-1, 2), good):
     59                     if not good_flag:
     60                         continue
     61                     tr.append((x, y))
     62                     if len(tr) > self.track_len:
     63                         del tr[0]
     64                     new_tracks.append(tr)
     65                     cv2.circle(vis, (x, y), 2, (0, 255, 0), -1)
     66                 self.tracks = new_tracks
     67                 cv2.polylines(vis, [np.int32(tr) for tr in self.tracks], False, (0, 255, 0))
     68                 draw_str(vis, (20, 20), 'track count: %d' % len(self.tracks))
     69 
     70             if self.frame_idx % self.detect_interval == 0:
     71                 mask = np.zeros_like(frame_gray)
     72                 mask[:] = 255
     73                 for x, y in [np.int32(tr[-1]) for tr in self.tracks]:
     74                     cv2.circle(mask, (x, y), 5, 0, -1)
     75                 p = cv2.goodFeaturesToTrack(frame_gray, mask = mask, **feature_params)
     76                 if p is not None:
     77                     for x, y in np.float32(p).reshape(-1, 2):
     78                         self.tracks.append([(x, y)])
     79 
     80 
     81             self.frame_idx += 1
     82             self.prev_gray = frame_gray
     83             cv2.imshow('lk_track', vis)
     84 
     85             ch = 0xFF & cv2.waitKey(1)
     86             if ch == 27:
     87                 break
     88 
     89 def main():
     90     import sys
     91     try:
     92         video_src = sys.argv[1]
     93     except:
     94         video_src = 0
     95 
     96     print __doc__
     97     App(video_src).run()
     98     cv2.destroyAllWindows()
     99 
    100 if __name__ == '__main__':
    101     main()
    102