1#!/usr/bin/env python 2 3''' 4example to show optical flow 5 6USAGE: opt_flow.py [<video_source>] 7 8Keys: 9 1 - toggle HSV flow visualization 10 2 - toggle glitch 11 12Keys: 13 ESC - exit 14''' 15 16# Python 2/3 compatibility 17from __future__ import print_function 18 19import numpy as np 20import cv2 as cv 21 22import video 23 24 25def draw_flow(img, flow, step=16): 26 h, w = img.shape[:2] 27 y, x = np.mgrid[step/2:h:step, step/2:w:step].reshape(2,-1).astype(int) 28 fx, fy = flow[y,x].T 29 lines = np.vstack([x, y, x+fx, y+fy]).T.reshape(-1, 2, 2) 30 lines = np.int32(lines + 0.5) 31 vis = cv.cvtColor(img, cv.COLOR_GRAY2BGR) 32 cv.polylines(vis, lines, 0, (0, 255, 0)) 33 for (x1, y1), (_x2, _y2) in lines: 34 cv.circle(vis, (x1, y1), 1, (0, 255, 0), -1) 35 return vis 36 37 38def draw_hsv(flow): 39 h, w = flow.shape[:2] 40 fx, fy = flow[:,:,0], flow[:,:,1] 41 ang = np.arctan2(fy, fx) + np.pi 42 v = np.sqrt(fx*fx+fy*fy) 43 hsv = np.zeros((h, w, 3), np.uint8) 44 hsv[...,0] = ang*(180/np.pi/2) 45 hsv[...,1] = 255 46 hsv[...,2] = np.minimum(v*4, 255) 47 bgr = cv.cvtColor(hsv, cv.COLOR_HSV2BGR) 48 return bgr 49 50 51def warp_flow(img, flow): 52 h, w = flow.shape[:2] 53 flow = -flow 54 flow[:,:,0] += np.arange(w) 55 flow[:,:,1] += np.arange(h)[:,np.newaxis] 56 res = cv.remap(img, flow, None, cv.INTER_LINEAR) 57 return res 58 59def main(): 60 import sys 61 try: 62 fn = sys.argv[1] 63 except IndexError: 64 fn = 0 65 66 cam = video.create_capture(fn) 67 _ret, prev = cam.read() 68 prevgray = cv.cvtColor(prev, cv.COLOR_BGR2GRAY) 69 show_hsv = False 70 show_glitch = False 71 cur_glitch = prev.copy() 72 73 while True: 74 _ret, img = cam.read() 75 gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY) 76 flow = cv.calcOpticalFlowFarneback(prevgray, gray, None, 0.5, 3, 15, 3, 5, 1.2, 0) 77 prevgray = gray 78 79 cv.imshow('flow', draw_flow(gray, flow)) 80 if show_hsv: 81 cv.imshow('flow HSV', draw_hsv(flow)) 82 if show_glitch: 83 cur_glitch = warp_flow(cur_glitch, flow) 84 cv.imshow('glitch', cur_glitch) 85 86 ch = cv.waitKey(5) 87 if ch == 27: 88 break 89 if ch == ord('1'): 90 show_hsv = not show_hsv 91 print('HSV flow visualization is', ['off', 'on'][show_hsv]) 92 if ch == ord('2'): 93 show_glitch = not show_glitch 94 if show_glitch: 95 cur_glitch = img.copy() 96 print('glitch is', ['off', 'on'][show_glitch]) 97 98 print('Done') 99 100 101if __name__ == '__main__': 102 print(__doc__) 103 main() 104 cv.destroyAllWindows() 105