Home | History | Annotate | Download | only in python2 
      1 #!/usr/bin/env python
      2 
      3 '''
      4 Lucas-Kanade homography 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. Finds homography between reference and current views.
     10 
     11 Usage
     12 -----
     13 lk_homography.py [<video_source>]
     14 
     15 
     16 Keys
     17 ----
     18 ESC   - exit
     19 SPACE - start tracking
     20 r     - toggle RANSAC
     21 '''
     22 
     23 import numpy as np
     24 import cv2
     25 import video
     26 from common import draw_str
     27 
     28 lk_params = dict( winSize  = (19, 19),
     29                   maxLevel = 2,
     30                   criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))
     31 
     32 feature_params = dict( maxCorners = 1000,
     33                        qualityLevel = 0.01,
     34                        minDistance = 8,
     35                        blockSize = 19 )
     36 
     37 def checkedTrace(img0, img1, p0, back_threshold = 1.0):
     38     p1, st, err = cv2.calcOpticalFlowPyrLK(img0, img1, p0, None, **lk_params)
     39     p0r, st, err = cv2.calcOpticalFlowPyrLK(img1, img0, p1, None, **lk_params)
     40     d = abs(p0-p0r).reshape(-1, 2).max(-1)
     41     status = d < back_threshold
     42     return p1, status
     43 
     44 green = (0, 255, 0)
     45 red = (0, 0, 255)
     46 
     47 class App:
     48     def __init__(self, video_src):
     49         self.cam = video.create_capture(video_src)
     50         self.p0 = None
     51         self.use_ransac = True
     52 
     53     def run(self):
     54         while True:
     55             ret, frame = self.cam.read()
     56             frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
     57             vis = frame.copy()
     58             if self.p0 is not None:
     59                 p2, trace_status = checkedTrace(self.gray1, frame_gray, self.p1)
     60 
     61                 self.p1 = p2[trace_status].copy()
     62                 self.p0 = self.p0[trace_status].copy()
     63                 self.gray1 = frame_gray
     64 
     65                 if len(self.p0) < 4:
     66                     self.p0 = None
     67                     continue
     68                 H, status = cv2.findHomography(self.p0, self.p1, (0, cv2.RANSAC)[self.use_ransac], 10.0)
     69                 h, w = frame.shape[:2]
     70                 overlay = cv2.warpPerspective(self.frame0, H, (w, h))
     71                 vis = cv2.addWeighted(vis, 0.5, overlay, 0.5, 0.0)
     72 
     73                 for (x0, y0), (x1, y1), good in zip(self.p0[:,0], self.p1[:,0], status[:,0]):
     74                     if good:
     75                         cv2.line(vis, (x0, y0), (x1, y1), (0, 128, 0))
     76                     cv2.circle(vis, (x1, y1), 2, (red, green)[good], -1)
     77                 draw_str(vis, (20, 20), 'track count: %d' % len(self.p1))
     78                 if self.use_ransac:
     79                     draw_str(vis, (20, 40), 'RANSAC')
     80             else:
     81                 p = cv2.goodFeaturesToTrack(frame_gray, **feature_params)
     82                 if p is not None:
     83                     for x, y in p[:,0]:
     84                         cv2.circle(vis, (x, y), 2, green, -1)
     85                     draw_str(vis, (20, 20), 'feature count: %d' % len(p))
     86 
     87             cv2.imshow('lk_homography', vis)
     88 
     89             ch = 0xFF & cv2.waitKey(1)
     90             if ch == 27:
     91                 break
     92             if ch == ord(' '):
     93                 self.frame0 = frame.copy()
     94                 self.p0 = cv2.goodFeaturesToTrack(frame_gray, **feature_params)
     95                 if self.p0 is not None:
     96                     self.p1 = self.p0
     97                     self.gray0 = frame_gray
     98                     self.gray1 = frame_gray
     99             if ch == ord('r'):
    100                 self.use_ransac = not self.use_ransac
    101 
    102 
    103 
    104 def main():
    105     import sys
    106     try:
    107         video_src = sys.argv[1]
    108     except:
    109         video_src = 0
    110 
    111     print __doc__
    112     App(video_src).run()
    113     cv2.destroyAllWindows()
    114 
    115 if __name__ == '__main__':
    116     main()
    117