This is a revision of mftrack.py by jayrambhia from here: https://github.com/jayrambhia/MFTracker

medianflowtracker.py

##############################################################################
# Full Imports
import time
import cv
import cv2
import sys

##############################################################################
# Partial Imports
from bb import getBB, getRectFromBB
from fbtrack import fbtrack as fbt #TODO Rename these

class MedianFlowTracker(object):
    
    def __init__(self, source = 0, boundingBox = None):
        '''
        source - The *full* path to a compatible video file, or an int that
                 specifies which camera the system should capture from.
        boundingBox - A 4-tuple specifying the x,y coords of two points. These
                      will be used to specify the bounding box.
        '''

        # If we are passed a video file, use that. Otherwise capture from the
        # first camera available to the system.
        self.source = source
        self.cam = cv2.VideoCapture(self.source)
        
        self.boundingBox = boundingBox

        # Mouse positions, start and end corners of the bounding box.
        self.mouseP1 = None
        self.mouseP2 = None
        
        # Is the mouse being dragged?
        self.mouseDrag = False

        # The image currently being processed
        self.img = None
        
        # If the mouse is being used
        self.mouseFlag = True
            
        # If we are not given a bounding box, let the user input one by
        # dragging their mouse over an image from the camera
        if not boundingBox:
            self.start()
        else:
            self.boundingBox = boundingBox
            _, self.img = self.cam.read()
            self.track()
        
    def start(self):
        '''
        Lets the user define an object to track by dragging a bounding box
        over an image from the given video stream
        '''
        
        # Grab one frame from the video stream
        _, self.img = self.cam.read()
        
        # Display that to the user
        cv2.imshow("img", self.img)
        
        # Pass our bounding box capture method to the OpenCV mouse handler
        cv.SetMouseCallback("img", self.__mouseHandler, None)
        
        if not self.boundingBox:
            _, self.img = self.cam.read()
            cv2.imshow("img", self.img)
            cv2.waitKey(30)
        
        cv2.waitKey(0)
    
    def __mouseHandler(self, event, x, y, flags, params):
        '''
        A method to be passed to the OpenCV mouse handler to capture the user
        dragging a bounding box on a presented image to define the object
        to be tracked.
        '''
        
        if self.mouseFlag == False:
            return
        
        # Get the first image from the camera
        _, self.img = self.cam.read()
        
        # The user has left-clicked for the first time. 
        if event == cv.CV_EVENT_LBUTTONDOWN and not self.mouseDrag:
            self.mouseP1 = (x, y)
            self.mouseDrag = True
            
        # The user has moved the mouse without letting up. Start drawing
        # a rectangle on the image.
        elif event == cv.CV_EVENT_MOUSEMOVE and self.mouseDrag:
            cv2.rectangle(self.img, self.mouseP1, (x, y), (255, 0, 0), 1, 8, 0)
            
        # The user has let up the mouse button after dragging. Capture the point
        # where they stopped. 
        elif event == cv.CV_EVENT_LBUTTONUP and self.mouseDrag:
            self.mouseP2 = (x, y)
            # Bounding box is too small
            if abs(self.mouseP1[0] - self.mouseP2[0]) < 10 or \
               abs(self.mouseP1[1] - self.mouseP2[1]) < 10:
                print 'Item too small. Please draw again.'
                self.mouseP1 = None
                self.mouseP2 = None
            self.mouseDrag=False
        
        # Display the latest image with overlayed rectangle.
        cv2.imshow("img",self.img)
        
        # Give the system time to do key capture
        cv2.waitKey(30)
        
        # Clean up
        if self.mouseP1 and self.mouseP2:
            cv2.destroyWindow("img")
            print self.mouseP1
            print self.mouseP2
            # Find the top left and bottom right corners
            xMax = max(self.mouseP1[0],self.mouseP2[0])
            xMin = min(self.mouseP1[0],self.mouseP2[0])
            yMax = max(self.mouseP1[1],self.mouseP2[1])
            yMin = min(self.mouseP1[1],self.mouseP2[1])
            self.boundingBox = [xMin, yMin, xMax - xMin, yMax - yMin]
            self.track()

    def track(self):
        '''
        Begin tracking the object defined by the bounding box
        '''
        self.mouseFlag = False
        
        # Convert our latest frame to grayscale
        oldGrayscaleImage = cv2.cvtColor(self.img, cv2.cv.CV_BGR2GRAY)
        
        # Extract the useful coords from our bounding box
        boundingBox = self.boundingBox
        boundingBox = [boundingBox[0], boundingBox[1],
                       boundingBox[0] + boundingBox[3],
                       boundingBox[1] + boundingBox[3]]
        
        # Main loop
        while True:

            t = time.time()
            # Get the latest frame
            _, img = self.cam.read()
            # Convert it to grayscale
            newGrayscaleImage = cv2.cvtColor(img, cv2.cv.CV_BGR2GRAY)
            
            # Calculate our shift between frames.
            newboundingBox, shift = fbt(oldGrayscaleImage,
                                        newGrayscaleImage,
                                        boundingBox, 12, 12, 3, 12)
            #print newboundingBox, "newboundingBox", shift, "shift"
            
            # Update frame and boundingbox info
            oldGrayscaleImage = newGrayscaleImage
            boundingBox = newboundingBox
            
            # Draw our bounding box on the latest image
            cv2.rectangle(img, (boundingBox[0], boundingBox[1]),
                               (boundingBox[2], boundingBox[3]),
                               (255, 0, 0))
            cv2.imshow("Media Flow Tracker", img)
            
            # Performance
            print (time.time() - t) * 1000, "msec"
            
            # Let the user quit by pressing the 'Esc' key
            k = cv2.waitKey(1)
            if k == 27:
                self.mouseFlag = False
                cv2.destroyAllWindows()
                # Simple return doesn't send signal to CV to exit.
                sys.exit()
            
        return