updated lt_track.py sample

- continuous feature detection
- track verification by reverse tracking

diff --git a/samples/python2/lk_track.py b/samples/python2/lk_track.py
index 0fe4cc6..c12598a 100644 (file)
--- a/samples/python2/lk_track.py
+++ b/samples/python2/lk_track.py
@@ -9,31 +9,71 @@ help_message = '''
 USAGE: lk_track.py [<video_source>]\r
 \r
 Keys:\r
-  1 - toggle old/new CalcOpticalFlowPyrLK implementation\r
   SPACE - reset features\r
 '''\r
 \r
-lk_params = dict( winSize  = (21, 21), \r
+lk_params = dict( winSize  = (15, 15), \r
                   maxLevel = 2, \r
                   criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03),\r
                   derivLambda = 0.0 )    \r
 \r
-feature_params = dict( maxCorners = 1000, \r
-                       qualityLevel = 0.1,\r
-                       minDistance = 5,\r
-                       blockSize = 5 )\r
+feature_params = dict( maxCorners = 500, \r
+                       qualityLevel = 0.3,\r
+                       minDistance = 7,\r
+                       blockSize = 7 )\r
 \r
-def calc_flow_old(img0, img1, p0):\r
-    p0 = [(x, y) for x, y in p0.reshape(-1, 2)]\r
-    h, w = img0.shape[:2]\r
-    img0_cv = cv.CreateMat(h, w, cv.CV_8U)\r
-    img1_cv = cv.CreateMat(h, w, cv.CV_8U)\r
-    np.asarray(img0_cv)[:] = img0\r
-    np.asarray(img1_cv)[:] = img1\r
-    t = clock()\r
-    features, status, error  = cv.CalcOpticalFlowPyrLK(img0_cv, img1_cv, None, None, p0, \r
-        lk_params['winSize'], lk_params['maxLevel'], lk_params['criteria'], 0, p0)\r
-    return np.float32(features), status, error, clock()-t\r
+class App:\r
+    def __init__(self, video_src):\r
+        self.track_len = 10\r
+        self.detect_interval = 5\r
+        self.tracks = []\r
+        self.cam = video.create_capture(video_src)\r
+        self.frame_idx = 0\r
+\r
+    def run(self):\r
+        while True:\r
+            ret, frame = self.cam.read()\r
+            frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)\r
+            vis = frame.copy()\r
+\r
+            if len(self.tracks) > 0:\r
+                img0, img1 = self.prev_gray, frame_gray\r
+                p0 = np.float32([tr[-1] for tr in self.tracks]).reshape(-1, 1, 2)\r
+                p1, st, err = cv2.calcOpticalFlowPyrLK(img0, img1, p0, None, **lk_params)\r
+                p0r, st, err = cv2.calcOpticalFlowPyrLK(img1, img0, p1, None, **lk_params)\r
+                d = abs(p0-p0r).reshape(-1, 2).max(-1)\r
+                good = d < 1\r
+                new_tracks = []\r
+                for tr, (x, y), good_flag in zip(self.tracks, p1.reshape(-1, 2), good):\r
+                    if not good_flag:\r
+                        continue\r
+                    tr.append((x, y))\r
+                    if len(tr) > self.track_len:\r
+                        del tr[0]\r
+                    new_tracks.append(tr)\r
+                    cv2.circle(vis, (x, y), 2, (0, 255, 0), -1)\r
+                self.tracks = new_tracks\r
+                cv2.polylines(vis, [np.int32(tr) for tr in self.tracks], False, (0, 255, 0))\r
+                draw_str(vis, (20, 20), 'track count: %d' % len(self.tracks))\r
+\r
+            if self.frame_idx % self.detect_interval == 0:\r
+                mask = np.zeros_like(frame_gray)\r
+                mask[:] = 255\r
+                for x, y in [np.int32(tr[-1]) for tr in self.tracks]:\r
+                    cv2.circle(mask, (x, y), 5, 0, -1)\r
+                p = cv2.goodFeaturesToTrack(frame_gray, mask = mask, **feature_params)\r
+                if p is not None:\r
+                    for x, y in np.float32(p).reshape(-1, 2):\r
+                        self.tracks.append([(x, y)])\r
+\r
+\r
+            self.frame_idx += 1\r
+            self.prev_gray = frame_gray\r
+            cv2.imshow('lk_track', vis)\r
+\r
+            ch = cv2.waitKey(1)\r
+            if ch == 27:\r
+                break\r
 \r
 def main():\r
     import sys\r
@@ -41,47 +81,7 @@ def main():
     except: video_src = video.presets['chess']\r
 \r
     print help_message\r
-\r
-    track_len = 4\r
-    tracks = []\r
-    cam = video.create_capture(video_src)\r
-    old_mode = True\r
-    while True:\r
-        ret, frame = cam.read()\r
-        vis = frame.copy()\r
-        if len(tracks) > 0:\r
-            p0 = np.float32([tr[-1] for tr in tracks]).reshape(-1, 1, 2)\r
-            img0 = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY)\r
-            img1 = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)\r
-            if old_mode:\r
-                p1,  st, err, dt = calc_flow_old(img0, img1, p0)\r
-            else:\r
-                t = clock()\r
-                p1,  st, err = cv2.calcOpticalFlowPyrLK(img0, img1, p0, None, **lk_params)\r
-                dt = clock()-t\r
-            for tr, (x, y) in zip(tracks, p1.reshape(-1, 2)):\r
-                tr.append((x, y))\r
-                if len(tr) > 10:\r
-                    del tr[0]\r
-                cv2.circle(vis, (x, y), 2, (0, 255, 0), -1)\r
-            cv2.polylines(vis, [np.int32(tr) for tr in tracks], False, (0, 255, 0))\r
-            draw_str(vis, (20, 20), ['new', 'old'][old_mode]+' mode')\r
-            draw_str(vis, (20, 40), 'time: %.02f ms' % (dt*1000))\r
-        prev_frame = frame.copy()\r
-\r
-        cv2.imshow('lk_track', vis)\r
-        ch = cv2.waitKey(5)\r
-        if ch == 27:\r
-            break\r
-        if ch == ord(' ') or len(tracks) == 0:\r
-            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)\r
-            p = cv2.goodFeaturesToTrack(gray, **feature_params)\r
-            p = [] if p is None else p.reshape(-1, 2)\r
-            tracks = []\r
-            for x, y in np.float32(p):\r
-                tracks.append([(x, y)])\r
-        if ch == ord('1'):\r
-            old_mode = not old_mode\r
+    App(video_src).run()\r
 \r
 if __name__ == '__main__':\r
     main()\r