added lk_track.py sample (currently as a comparison between flow tracker implementations)

diff --git a/samples/python2/facedetect.py b/samples/python2/facedetect.py
index 4eed36f..ca85b14 100644 (file)
--- a/samples/python2/facedetect.py
+++ b/samples/python2/facedetect.py
@@ -15,7 +15,7 @@ def detect(img, cascade):
 \r
 def draw_rects(img, rects, color):\r
     for x1, y1, x2, y2 in rects:\r
-        cv2.rectangle(img, (x1, y1), (x2, y2), color)\r
+        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)\r
 \r
 if __name__ == '__main__':\r
     import sys, getopt\r
@@ -38,7 +38,7 @@ if __name__ == '__main__':
         gray = cv2.cvtColor(img, cv.CV_BGR2GRAY)\r
         gray = cv2.equalizeHist(gray)\r
         rects = detect(gray, cascade)\r
-        vis = cv2.cvtColor(gray, cv.CV_GRAY2BGR)\r
+        vis = img.copy()\r
         draw_rects(vis, rects, (0, 255, 0))\r
         for x1, y1, x2, y2 in rects:\r
             roi = gray[y1:y2, x1:x2]\r

diff --git a/samples/python2/lk_track.py b/samples/python2/lk_track.py
new file mode 100644 (file)
index 0000000..65346de
--- /dev/null
+++ b/samples/python2/lk_track.py
@@ -0,0 +1,87 @@
+import numpy as np\r
+import cv2, cv\r
+import video\r
+from common import anorm2, draw_str\r
+from time import clock\r
+\r
+help_message = '''\r
+USAGE: lk_track.py [<video_source>]\r
+\r
+Keys:\r
+  1 - toggle old/new CalcOpticalFlowPyrLK implementation\r
+  SPACE - reset features\r
+'''\r
+\r
+\r
+\r
+\r
+lk_params = dict( winSize  = (3, 3), \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 = 500, \r
+                       qualityLevel = 0.1,\r
+                       minDistance = 5,\r
+                       blockSize = 5 )\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'], (cv.CV_TERMCRIT_EPS | cv.CV_TERMCRIT_ITER, 10, 0.03), 0, p0)\r
+    return np.float32(features), status, error, clock()-t\r
+\r
+def main():\r
+    import sys\r
+    try: video_src = sys.argv[1]\r
+    except: video_src = video.presets['chess']\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, cv.CV_BGR2GRAY)\r
+            img1 = cv2.cvtColor(frame, cv.CV_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, **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), -2)\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('img', 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, cv.CV_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
+\r
+if __name__ == '__main__':\r
+    main()\r