#! /usr/bin/env python
# -*- coding: utf-8 -*-

import cv
import opticalflow



# Open the camera up...
vid = cv.CaptureFromCAM(-1)

width = int(cv.GetCaptureProperty(vid,cv.CV_CAP_PROP_FRAME_WIDTH))
height = int(cv.GetCaptureProperty(vid,cv.CV_CAP_PROP_FRAME_HEIGHT))
fps = cv.GetCaptureProperty(vid,cv.CV_CAP_PROP_FPS)



# Setup windows...
windowInput = 'Input'
cv.NamedWindow(windowInput)
cv.ResizeWindow(windowInput,width,height)

windowOptical = 'Optical Flow'
cv.NamedWindow(windowOptical)
cv.ResizeWindow(windowOptical,width,height)



# Buffers for the optical flow result...
size = (width,height)
mask = cv.CreateImage(size,cv.IPL_DEPTH_32F,1)
velX = cv.CreateImage(size,cv.IPL_DEPTH_32F,1)
velY = cv.CreateImage(size,cv.IPL_DEPTH_32F,1)
velMult = 0.2

display = cv.CreateImage(size,cv.IPL_DEPTH_32F,3)



# Optical flow generation object...
oflow = opticalflow.OpticalFlow(width,height)
oflow.updateParam(11,2.0,2.0,1,0)



# Eat frames and process...
while True:
  # Get the frame...
  frame = cv.QueryFrame(vid)
  if frame==None:
    key = cv.WaitKey(5)
    if key==27:
      break
    else:
      continue
  cv.ShowImage(windowInput,frame)

  # Do optical flow...
  oflow.nextFrame(frame)

  # Get the optical flow output...
  cv.Copy(oflow.getU(),velX)
  cv.Copy(oflow.getV(),velY)

  # Process and display the optical flow output...
  cv.ConvertScale(velX,velX,velMult,0.5)
  cv.ConvertScale(velY,velY,velMult,0.5)
  cv.ConvertScale(oflow.getMask(), mask, 1.0, 0.0)
  cv.Mul(velX,mask,velX,1.0)
  cv.Mul(velY,mask,velY,1.0)
  cv.Zero(display)
  cv.MixChannels([velX],[display],[(0,2)])
  cv.MixChannels([velY],[display],[(0,1)])

  # Get and render the top maxima...
  maxima = oflow.getMaxima()
  cv.Circle(display,(maxima[0,0],maxima[0,1]),5,(1,0,0))

  # Display the optical flow...
  cv.ShowImage(windowOptical,display)
  

  # Check if we need to exit...
  key = cv.WaitKey(5)
  if key==27:
    break