From edf0b40d42bf887167a2958ef254796594f6fc21 Mon Sep 17 00:00:00 2001
From: Alexander Mordvintsev <no@email>
Date: Thu, 7 Jul 2011 07:55:49 +0000
Subject: [PATCH] opt_flow.py demo

---
 samples/python2/opt_flow.py | 97 +++++++++++++++++++++++++++++++++------------
 samples/python2/video.py    |  4 +-
 2 files changed, 74 insertions(+), 27 deletions(-)

diff --git a/samples/python2/opt_flow.py b/samples/python2/opt_flow.py
index 1c30294..4a18664 100644
--- a/samples/python2/opt_flow.py
+++ b/samples/python2/opt_flow.py
@@ -1,38 +1,83 @@
 import numpy as np
 import cv2, cv
 import video
-import sys
 
-try: fn = sys.argv[1]
-except: fn = video.presets['chess']
+help_message = '''
+USAGE: opt_flow.py [<video_source>]
 
-cam = video.create_capture(fn)
-ret, prev = cam.read()
-#prev = cv2.pyrDown(prev)
-prevgray = cv2.cvtColor(prev, cv.CV_BGR2GRAY)
+Keys:
+ 1 - toggle HSV flow visualization
+ 2 - toggle glitch
 
-def draw_flow(img, flow, step):
+'''
+
+def draw_flow(img, flow, step=16):
     h, w = img.shape[:2]
-    y, x = map(np.ravel, np.mgrid[step/2:h:step, step/2:w:step])
-    f = flow[y,x]
-    x1 = x + f[:,0]
-    y1 = y + f[:,1]
-    #lines = np.int32( np.vstack([x, y, x1, y1]).T )
+    y, x = np.mgrid[step/2:h:step, step/2:w:step].reshape(2,-1)
+    fx, fy = flow[y,x].T
+    lines = np.vstack([x, y, x+fx, y+fy]).T.reshape(-1, 2, 2)
+    lines = np.int32(lines + 0.5)
     vis = cv2.cvtColor(img, cv.CV_GRAY2BGR)
-    #print lines
-    #cv2.polylines(vis, lines, 0, (0, 255, 0))
-    for x_, y_, x1_, y1_ in np.int32(zip(x, y, x1, y1)):
-        cv2.line(vis, (x_, y_), (x1_, y1_), (0, 255, 0))
+    cv2.polylines(vis, lines, 0, (0, 255, 0))
+    for (x1, y1), (x2, y2) in lines:
+        cv2.circle(vis, (x1, y1), 1, (0, 255, 0), -1)
     return vis
 
-while True:
-    ret, img = cam.read()
-    #img = cv2.pyrDown(img)
-    gray = cv2.cvtColor(img, cv.CV_BGR2GRAY)
-    flow = cv2.calcOpticalFlowFarneback(prevgray, gray, 0.5, 3, 15, 3, 5, 1.2, 0)
-    prevgray = gray
+def draw_hsv(flow):
+    h, w = flow.shape[:2]
+    fx, fy = flow[:,:,0], flow[:,:,1]
+    ang = np.arctan2(fy, fx) + np.pi
+    v = np.sqrt(fx*fx+fy*fy)
+    hsv = np.zeros((h, w, 3), np.uint8)
+    hsv[...,0] = np.rad2deg(ang)/2
+    hsv[...,1] = 255
+    hsv[...,2] = np.minimum(v*4, 255)
+    bgr = cv2.cvtColor(hsv, cv.CV_HSV2BGR)
+    return bgr
+
+def warp_flow(img, flow):
+    h, w = flow.shape[:2]
+    flow = -flow
+    flow[:,:,0] += np.arange(w)
+    flow[:,:,1] += np.arange(h)[:,np.newaxis]
+    res = cv2.remap(img, flow, None, cv2.INTER_LINEAR)
+    return res
+
+if __name__ == '__main__':
+    import sys
+    print help_message
+    try: fn = sys.argv[1]
+    except: fn = video.presets['chess']
+
+    cam = video.create_capture(fn)
+    ret, prev = cam.read()
+    prevgray = cv2.cvtColor(prev, cv.CV_BGR2GRAY)
+    show_hsv = False
+    show_glitch = False
+    cur_glitch = prev.copy()
+
+    while True:
+        ret, img = cam.read()
+        gray = cv2.cvtColor(img, cv.CV_BGR2GRAY)
+        flow = cv2.calcOpticalFlowFarneback(prevgray, gray, 0.5, 3, 15, 3, 5, 1.2, 0)
+        prevgray = gray
+        
+        cv2.imshow('flow', draw_flow(gray, flow))
+        if show_hsv:
+            cv2.imshow('flow HSV', draw_hsv(flow))
+        if show_glitch:
+            cur_glitch = warp_flow(cur_glitch, flow)
+            cv2.imshow('glitch', cur_glitch)
 
-    cv2.imshow('flow', draw_flow(gray, flow, 16))
-    if cv2.waitKey(5) == 27:
-        break
+        ch = cv2.waitKey(5)
+        if ch == 27:
+            break
+        if ch == ord('1'):
+            show_hsv = not show_hsv
+            print 'HSV flow visualization is', ['off', 'on'][show_hsv]
+        if ch == ord('2'):
+            show_glitch = not show_glitch
+            if show_glitch:
+                cur_glitch = img.copy()
+            print 'glitch is', ['off', 'on'][show_glitch]
 
diff --git a/samples/python2/video.py b/samples/python2/video.py
index 12bc931..d18ee82 100644
--- a/samples/python2/video.py
+++ b/samples/python2/video.py
@@ -60,6 +60,7 @@ class Chess(VideoSynthBase):
                         [0.0,0.0,      1.0]])
 
         self.dist_coef = np.float64([-0.2, 0.1, 0, 0])
+        self.t = 0
 
     def draw_quads(self, img, quads, color = (0, 255, 0)):
         img_quads = cv2.projectPoints(quads.reshape(-1, 3), self.rvec, self.tvec, self.K, self.dist_coef) [0]
@@ -68,7 +69,8 @@ class Chess(VideoSynthBase):
             cv2.fillConvexPoly(img, np.int32(q*4), color, cv2.CV_AA, shift=2)
 
     def render(self, dst):
-        t = clock()
+        t = self.t
+        self.t += 1.0/30.0
         
         sx, sy = self.grid_size
         center = np.array([0.5*sx, 0.5*sy, 0.0])
-- 
2.7.4