• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python
2
3'''
4Lucas-Kanade tracker
5====================
6
7Lucas-Kanade sparse optical flow demo. Uses goodFeaturesToTrack
8for track initialization and back-tracking for match verification
9between frames.
10
11Usage
12-----
13lk_track.py [<video_source>]
14
15
16Keys
17----
18ESC - exit
19'''
20
21import numpy as np
22import cv2
23import video
24from common import anorm2, draw_str
25from time import clock
26
27lk_params = dict( winSize  = (15, 15),
28                  maxLevel = 2,
29                  criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))
30
31feature_params = dict( maxCorners = 500,
32                       qualityLevel = 0.3,
33                       minDistance = 7,
34                       blockSize = 7 )
35
36class 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
89def 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
100if __name__ == '__main__':
101    main()
102