#!/usr/bin/env python from __future__ import print_function import cv2, os, sys, time import numpy as np from argparse import ArgumentParser def draw(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) f = np.zeros((h, w, 3), np.uint8) f[...,0] = 0 #np.minimum(v*10, 255) f[...,1] = 0 f[...,2] = 255- np.minimum(v**3, 255) #ang*(180/np.pi/2) bgr = cv2.cvtColor(f, cv2.COLOR_HSV2BGR) return bgr p = ArgumentParser("") p.add_argument("--video", type=int, default=0, help="video, default: 0") p.add_argument("--output", default=None, help="path to save movie, default: None (show live)") p.add_argument("--width", type=int, default=640, help="pre-detect resize width") p.add_argument("--height", type=int, default=480, help="pre-detect resize height") p.add_argument("--fourcc", default="XVID", help="MJPG,mp4v,XVID") p.add_argument("--framerate", type=float, default=25, help="output frame rate") p.add_argument("--show", default=False, action="store_true") p.add_argument("--frames", type=int, default=100) args = p.parse_args() fourcc = None cam = cv2.VideoCapture(args.video) cam.set(cv2.cv.CV_CAP_PROP_FRAME_WIDTH, args.width) cam.set(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT, args.height) if args.output: try: fourcc = cv2.cv.CV_FOURCC(*args.fourcc) except AttributeError: fourcc = cv2.VideoWriter_fourcc(*args.fourcc) out = cv2.VideoWriter() out.open(args.output, fourcc, args.framerate, (args.width, args.height)) else: out = None while True: ret, prev = cam.read() prevgray = cv2.cvtColor(prev, cv2.COLOR_BGR2GRAY) if prevgray.shape == (args.height, args.width): break count = 0 try: while True: ret, frame = cam.read() gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) flow = cv2.calcOpticalFlowFarneback(prevgray, gray, 0.5, 3, 15, 3, 5, 1.2, 0) prevgray = gray frame = draw(flow) if out != None: out.write(frame) count += 1 if args.show: cv2.imshow('display', frame) if cv2.waitKey(5) & 0xFF == ord('q'): break if args.frames != None: if (count >= args.frames): break except KeyboardInterrupt: pass print ("\nCleaning up... Wrote", count, "frames") if out: out.release() if args.show: cv2.destroyAllWindows()