opt_flow.py demo

diff --git a/samples/python2/opt_flow.py b/samples/python2/opt_flow.py
index 1c30294..4a18664 100644 (file)
--- a/samples/python2/opt_flow.py
+++ b/samples/python2/opt_flow.py
@@ -1,38 +1,83 @@
 import numpy as np\r
 import cv2, cv\r
 import video\r
-import sys\r
 \r
-try: fn = sys.argv[1]\r
-except: fn = video.presets['chess']\r
+help_message = '''\r
+USAGE: opt_flow.py [<video_source>]\r
 \r
-cam = video.create_capture(fn)\r
-ret, prev = cam.read()\r
-#prev = cv2.pyrDown(prev)\r
-prevgray = cv2.cvtColor(prev, cv.CV_BGR2GRAY)\r
+Keys:\r
+ 1 - toggle HSV flow visualization\r
+ 2 - toggle glitch\r
 \r
-def draw_flow(img, flow, step):\r
+'''\r
+\r
+def draw_flow(img, flow, step=16):\r
     h, w = img.shape[:2]\r
-    y, x = map(np.ravel, np.mgrid[step/2:h:step, step/2:w:step])\r
-    f = flow[y,x]\r
-    x1 = x + f[:,0]\r
-    y1 = y + f[:,1]\r
-    #lines = np.int32( np.vstack([x, y, x1, y1]).T )\r
+    y, x = np.mgrid[step/2:h:step, step/2:w:step].reshape(2,-1)\r
+    fx, fy = flow[y,x].T\r
+    lines = np.vstack([x, y, x+fx, y+fy]).T.reshape(-1, 2, 2)\r
+    lines = np.int32(lines + 0.5)\r
     vis = cv2.cvtColor(img, cv.CV_GRAY2BGR)\r
-    #print lines\r
-    #cv2.polylines(vis, lines, 0, (0, 255, 0))\r
-    for x_, y_, x1_, y1_ in np.int32(zip(x, y, x1, y1)):\r
-        cv2.line(vis, (x_, y_), (x1_, y1_), (0, 255, 0))\r
+    cv2.polylines(vis, lines, 0, (0, 255, 0))\r
+    for (x1, y1), (x2, y2) in lines:\r
+        cv2.circle(vis, (x1, y1), 1, (0, 255, 0), -1)\r
     return vis\r
 \r
-while True:\r
-    ret, img = cam.read()\r
-    #img = cv2.pyrDown(img)\r
-    gray = cv2.cvtColor(img, cv.CV_BGR2GRAY)\r
-    flow = cv2.calcOpticalFlowFarneback(prevgray, gray, 0.5, 3, 15, 3, 5, 1.2, 0)\r
-    prevgray = gray\r
+def draw_hsv(flow):\r
+    h, w = flow.shape[:2]\r
+    fx, fy = flow[:,:,0], flow[:,:,1]\r
+    ang = np.arctan2(fy, fx) + np.pi\r
+    v = np.sqrt(fx*fx+fy*fy)\r
+    hsv = np.zeros((h, w, 3), np.uint8)\r
+    hsv[...,0] = np.rad2deg(ang)/2\r
+    hsv[...,1] = 255\r
+    hsv[...,2] = np.minimum(v*4, 255)\r
+    bgr = cv2.cvtColor(hsv, cv.CV_HSV2BGR)\r
+    return bgr\r
+\r
+def warp_flow(img, flow):\r
+    h, w = flow.shape[:2]\r
+    flow = -flow\r
+    flow[:,:,0] += np.arange(w)\r
+    flow[:,:,1] += np.arange(h)[:,np.newaxis]\r
+    res = cv2.remap(img, flow, None, cv2.INTER_LINEAR)\r
+    return res\r
+\r
+if __name__ == '__main__':\r
+    import sys\r
+    print help_message\r
+    try: fn = sys.argv[1]\r
+    except: fn = video.presets['chess']\r
+\r
+    cam = video.create_capture(fn)\r
+    ret, prev = cam.read()\r
+    prevgray = cv2.cvtColor(prev, cv.CV_BGR2GRAY)\r
+    show_hsv = False\r
+    show_glitch = False\r
+    cur_glitch = prev.copy()\r
+\r
+    while True:\r
+        ret, img = cam.read()\r
+        gray = cv2.cvtColor(img, cv.CV_BGR2GRAY)\r
+        flow = cv2.calcOpticalFlowFarneback(prevgray, gray, 0.5, 3, 15, 3, 5, 1.2, 0)\r
+        prevgray = gray\r
+        \r
+        cv2.imshow('flow', draw_flow(gray, flow))\r
+        if show_hsv:\r
+            cv2.imshow('flow HSV', draw_hsv(flow))\r
+        if show_glitch:\r
+            cur_glitch = warp_flow(cur_glitch, flow)\r
+            cv2.imshow('glitch', cur_glitch)\r
 \r
-    cv2.imshow('flow', draw_flow(gray, flow, 16))\r
-    if cv2.waitKey(5) == 27:\r
-        break\r
+        ch = cv2.waitKey(5)\r
+        if ch == 27:\r
+            break\r
+        if ch == ord('1'):\r
+            show_hsv = not show_hsv\r
+            print 'HSV flow visualization is', ['off', 'on'][show_hsv]\r
+        if ch == ord('2'):\r
+            show_glitch = not show_glitch\r
+            if show_glitch:\r
+                cur_glitch = img.copy()\r
+            print 'glitch is', ['off', 'on'][show_glitch]\r
 \r

diff --git a/samples/python2/video.py b/samples/python2/video.py
index 12bc931..d18ee82 100644 (file)
--- a/samples/python2/video.py
+++ b/samples/python2/video.py
@@ -60,6 +60,7 @@ class Chess(VideoSynthBase):
                         [0.0,0.0,      1.0]])\r
 \r
         self.dist_coef = np.float64([-0.2, 0.1, 0, 0])\r
+        self.t = 0\r
 \r
     def draw_quads(self, img, quads, color = (0, 255, 0)):\r
         img_quads = cv2.projectPoints(quads.reshape(-1, 3), self.rvec, self.tvec, self.K, self.dist_coef) [0]\r
@@ -68,7 +69,8 @@ class Chess(VideoSynthBase):
             cv2.fillConvexPoly(img, np.int32(q*4), color, cv2.CV_AA, shift=2)\r
 \r
     def render(self, dst):\r
-        t = clock()\r
+        t = self.t\r
+        self.t += 1.0/30.0\r
         \r
         sx, sy = self.grid_size\r
         center = np.array([0.5*sx, 0.5*sy, 0.0])\r