Home | History | Annotate | Download | only in python2
      1 #!/usr/bin/env python
      2 
      3 '''
      4 Camshift tracker
      5 ================
      6 
      7 This is a demo that shows mean-shift based tracking
      8 You select a color objects such as your face and it tracks it.
      9 This reads from video camera (0 by default, or the camera number the user enters)
     10 
     11 http://www.robinhewitt.com/research/track/camshift.html
     12 
     13 Usage:
     14 ------
     15     camshift.py [<video source>]
     16 
     17     To initialize tracking, select the object with mouse
     18 
     19 Keys:
     20 -----
     21     ESC   - exit
     22     b     - toggle back-projected probability visualization
     23 '''
     24 
     25 import numpy as np
     26 import cv2
     27 
     28 # local module
     29 import video
     30 
     31 
     32 class App(object):
     33     def __init__(self, video_src):
     34         self.cam = video.create_capture(video_src)
     35         ret, self.frame = self.cam.read()
     36         cv2.namedWindow('camshift')
     37         cv2.setMouseCallback('camshift', self.onmouse)
     38 
     39         self.selection = None
     40         self.drag_start = None
     41         self.tracking_state = 0
     42         self.show_backproj = False
     43 
     44     def onmouse(self, event, x, y, flags, param):
     45         x, y = np.int16([x, y]) # BUG
     46         if event == cv2.EVENT_LBUTTONDOWN:
     47             self.drag_start = (x, y)
     48             self.tracking_state = 0
     49         if self.drag_start:
     50             if flags & cv2.EVENT_FLAG_LBUTTON:
     51                 h, w = self.frame.shape[:2]
     52                 xo, yo = self.drag_start
     53                 x0, y0 = np.maximum(0, np.minimum([xo, yo], [x, y]))
     54                 x1, y1 = np.minimum([w, h], np.maximum([xo, yo], [x, y]))
     55                 self.selection = None
     56                 if x1-x0 > 0 and y1-y0 > 0:
     57                     self.selection = (x0, y0, x1, y1)
     58             else:
     59                 self.drag_start = None
     60                 if self.selection is not None:
     61                     self.tracking_state = 1
     62 
     63     def show_hist(self):
     64         bin_count = self.hist.shape[0]
     65         bin_w = 24
     66         img = np.zeros((256, bin_count*bin_w, 3), np.uint8)
     67         for i in xrange(bin_count):
     68             h = int(self.hist[i])
     69             cv2.rectangle(img, (i*bin_w+2, 255), ((i+1)*bin_w-2, 255-h), (int(180.0*i/bin_count), 255, 255), -1)
     70         img = cv2.cvtColor(img, cv2.COLOR_HSV2BGR)
     71         cv2.imshow('hist', img)
     72 
     73     def run(self):
     74         while True:
     75             ret, self.frame = self.cam.read()
     76             vis = self.frame.copy()
     77             hsv = cv2.cvtColor(self.frame, cv2.COLOR_BGR2HSV)
     78             mask = cv2.inRange(hsv, np.array((0., 60., 32.)), np.array((180., 255., 255.)))
     79 
     80             if self.selection:
     81                 x0, y0, x1, y1 = self.selection
     82                 self.track_window = (x0, y0, x1-x0, y1-y0)
     83                 hsv_roi = hsv[y0:y1, x0:x1]
     84                 mask_roi = mask[y0:y1, x0:x1]
     85                 hist = cv2.calcHist( [hsv_roi], [0], mask_roi, [16], [0, 180] )
     86                 cv2.normalize(hist, hist, 0, 255, cv2.NORM_MINMAX);
     87                 self.hist = hist.reshape(-1)
     88                 self.show_hist()
     89 
     90                 vis_roi = vis[y0:y1, x0:x1]
     91                 cv2.bitwise_not(vis_roi, vis_roi)
     92                 vis[mask == 0] = 0
     93 
     94             if self.tracking_state == 1:
     95                 self.selection = None
     96                 prob = cv2.calcBackProject([hsv], [0], self.hist, [0, 180], 1)
     97                 prob &= mask
     98                 term_crit = ( cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 1 )
     99                 track_box, self.track_window = cv2.CamShift(prob, self.track_window, term_crit)
    100 
    101                 if self.show_backproj:
    102                     vis[:] = prob[...,np.newaxis]
    103                 try:
    104                     cv2.ellipse(vis, track_box, (0, 0, 255), 2)
    105                 except:
    106                     print track_box
    107 
    108             cv2.imshow('camshift', vis)
    109 
    110             ch = 0xFF & cv2.waitKey(5)
    111             if ch == 27:
    112                 break
    113             if ch == ord('b'):
    114                 self.show_backproj = not self.show_backproj
    115         cv2.destroyAllWindows()
    116 
    117 
    118 if __name__ == '__main__':
    119     import sys
    120     try:
    121         video_src = sys.argv[1]
    122     except:
    123         video_src = 0
    124     print __doc__
    125     App(video_src).run()
    126