opt_flow.py 2.33 KB
Newer Older
wester committed
1 2
#!/usr/bin/env python

wester committed
3 4 5
import numpy as np
import cv2
import video
wester committed
6

wester committed
7
help_message = '''
wester committed
8 9 10 11 12 13 14 15 16 17
USAGE: opt_flow.py [<video_source>]

Keys:
 1 - toggle HSV flow visualization
 2 - toggle glitch

'''

def draw_flow(img, flow, step=16):
    h, w = img.shape[:2]
wester committed
18
    y, x = np.mgrid[step/2:h:step, step/2:w:step].reshape(2,-1)
wester committed
19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49
    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, cv2.COLOR_GRAY2BGR)
    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

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] = ang*(180/np.pi/2)
    hsv[...,1] = 255
    hsv[...,2] = np.minimum(v*4, 255)
    bgr = cv2.cvtColor(hsv, cv2.COLOR_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
wester committed
50 51 52
    print help_message
    try: fn = sys.argv[1]
    except: fn = 0
wester committed
53 54 55 56 57 58 59 60 61 62 63

    cam = video.create_capture(fn)
    ret, prev = cam.read()
    prevgray = cv2.cvtColor(prev, cv2.COLOR_BGR2GRAY)
    show_hsv = False
    show_glitch = False
    cur_glitch = prev.copy()

    while True:
        ret, img = cam.read()
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
wester committed
64
        flow = cv2.calcOpticalFlowFarneback(prevgray, gray, 0.5, 3, 15, 3, 5, 1.2, 0)
wester committed
65 66 67 68 69 70 71 72 73
        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)

a  
Kai Westerkamp committed
74
        ch = 0xFF & cv2.waitKey(5)
wester committed
75 76 77 78
        if ch == 27:
            break
        if ch == ord('1'):
            show_hsv = not show_hsv
wester committed
79
            print 'HSV flow visualization is', ['off', 'on'][show_hsv]
wester committed
80 81 82 83
        if ch == ord('2'):
            show_glitch = not show_glitch
            if show_glitch:
                cur_glitch = img.copy()
wester committed
84
            print 'glitch is', ['off', 'on'][show_glitch]
wester committed
85
    cv2.destroyAllWindows()