git.maemo.org Git - opencv/blob - samples/python/lkdemo.py

   1 #! /usr/bin/env python
   2
   3 print "OpenCV Python version of lkdemo"
   4
   5 import sys
   6
   7 # import the necessary things for OpenCV
   8 from opencv import cv
   9 from opencv import highgui
  10
  11 #############################################################################
  12 # some "constants"
  13
  14 win_size = 10
  15 MAX_COUNT = 500
  16
  17 #############################################################################
  18 # some "global" variables
  19
  20 image = None
  21 pt = None
  22 add_remove_pt = False
  23 flags = 0
  24 night_mode = False
  25 need_to_init = False
  26
  27 #############################################################################
  28 # the mouse callback
  29
  30 # the callback on the trackbar
  31 def on_mouse (event, x, y, flags, param):
  32
  33     # we will use the global pt and add_remove_pt
  34     global pt
  35     global add_remove_pt
  36
  37     if image is None:
  38         # not initialized, so skip
  39         return
  40
  41     if image.origin != 0:
  42         # different origin
  43         y = image.height - y
  44
  45     if event == highgui.CV_EVENT_LBUTTONDOWN:
  46         # user has click, so memorize it
  47         pt = cv.cvPoint (x, y)
  48         add_remove_pt = True
  49
  50 #############################################################################
  51 # so, here is the main part of the program
  52
  53 if __name__ == '__main__':
  54
  55     try:
  56         # try to get the device number from the command line
  57         device = int (sys.argv [1])
  58
  59         # got it ! so remove it from the arguments
  60         del sys.argv [1]
  61     except (IndexError, ValueError):
  62         # no device number on the command line, assume we want the 1st device
  63         device = 0
  64
  65     if len (sys.argv) == 1:
  66         # no argument on the command line, try to use the camera
  67         capture = highgui.cvCreateCameraCapture (device)
  68
  69     else:
  70         # we have an argument on the command line,
  71         # we can assume this is a file name, so open it
  72         capture = highgui.cvCreateFileCapture (sys.argv [1])
  73
  74     # check that capture device is OK
  75     if not capture:
  76         print "Error opening capture device"
  77         sys.exit (1)
  78
  79     # display a small howto use it
  80     print "Hot keys: \n" \
  81           "\tESC - quit the program\n" \
  82           "\tr - auto-initialize tracking\n" \
  83           "\tc - delete all the points\n" \
  84           "\tn - switch the \"night\" mode on/off\n" \
  85           "To add/remove a feature point click it\n"
  86
  87     # first, create the necessary windows
  88     highgui.cvNamedWindow ('LkDemo', highgui.CV_WINDOW_AUTOSIZE)
  89
  90     # register the mouse callback
  91     highgui.cvSetMouseCallback ('LkDemo', on_mouse, None)
  92
  93     while 1:
  94         # do forever
  95
  96         # 1. capture the current image
  97         frame = highgui.cvQueryFrame (capture)
  98         if frame is None:
  99             # no image captured... end the processing
 100             break
 101
 102         if image is None:
 103             # create the images we need
 104             image = cv.cvCreateImage (cv.cvGetSize (frame), 8, 3)
 105             image.origin = frame.origin
 106             grey = cv.cvCreateImage (cv.cvGetSize (frame), 8, 1)
 107             prev_grey = cv.cvCreateImage (cv.cvGetSize (frame), 8, 1)
 108             pyramid = cv.cvCreateImage (cv.cvGetSize (frame), 8, 1)
 109             prev_pyramid = cv.cvCreateImage (cv.cvGetSize (frame), 8, 1)
 110             points = [[], []]
 111
 112         # copy the frame, so we can draw on it
 113         cv.cvCopy (frame, image)
 114
 115         # create a grey version of the image
 116         cv.cvCvtColor (image, grey, cv.CV_BGR2GRAY)
 117
 118         if night_mode:
 119             # night mode: only display the points
 120             cv.cvSetZero (image)
 121
 122         if need_to_init:
 123             # we want to search all the good points
 124
 125             # create the wanted images
 126             eig = cv.cvCreateImage (cv.cvGetSize (grey), 32, 1)
 127             temp = cv.cvCreateImage (cv.cvGetSize (grey), 32, 1)
 128
 129             # the default parameters
 130             quality = 0.01
 131             min_distance = 10
 132
 133             # search the good points
 134             points [1] = cv.cvGoodFeaturesToTrack (
 135                 grey, eig, temp,
 136                 MAX_COUNT,
 137                 quality, min_distance, None, 3, 0, 0.04)
 138
 139             # refine the corner locations
 140             cv.cvFindCornerSubPix (
 141                 grey,
 142                 points [1],
 143                 cv.cvSize (win_size, win_size), cv.cvSize (-1, -1),
 144                 cv.cvTermCriteria (cv.CV_TERMCRIT_ITER | cv.CV_TERMCRIT_EPS,
 145                                    20, 0.03))
 146
 147             # release the temporary images
 148             cv.cvReleaseImage (eig)
 149             cv.cvReleaseImage (temp)
 150
 151         elif len (points [0]) > 0:
 152             # we have points, so display them
 153
 154             # calculate the optical flow
 155             points [1], status = cv.cvCalcOpticalFlowPyrLK (
 156                 prev_grey, grey, prev_pyramid, pyramid,
 157                 points [0], len (points [0]),
 158                 cv.cvSize (win_size, win_size), 3,
 159                 len (points [0]),
 160                 None,
 161                 cv.cvTermCriteria (cv.CV_TERMCRIT_ITER|cv.CV_TERMCRIT_EPS,
 162                                    20, 0.03),
 163                 flags)
 164
 165             # initializations
 166             point_counter = -1
 167             new_points = []
 168
 169             for the_point in points [1]:
 170                 # go trough all the points
 171
 172                 # increment the counter
 173                 point_counter += 1
 174
 175                 if add_remove_pt:
 176                     # we have a point to add, so see if it is close to
 177                     # another one. If yes, don't use it
 178                     dx = pt.x - the_point.x
 179                     dy = pt.y - the_point.y
 180                     if dx * dx + dy * dy <= 25:
 181                         # too close
 182                         add_remove_pt = 0
 183                         continue
 184
 185                 if not status [point_counter]:
 186                     # we will disable this point
 187                     continue
 188
 189                 # this point is a correct point
 190                 new_points.append (the_point)
 191
 192                 # draw the current point
 193                 cv.cvCircle (image,
 194                              [the_point.x, the_point.y],
 195                              3, cv.cvScalar (0, 255, 0, 0),
 196                              -1, 8, 0)
 197
 198             # set back the points we keep
 199             points [1] = new_points
 200
 201         if add_remove_pt:
 202             # we want to add a point
 203             points [1].append (cv.cvPointTo32f (pt))
 204
 205             # refine the corner locations
 206             points [1][-1] = cv.cvFindCornerSubPix (
 207                 grey,
 208                 [points [1][-1]],
 209                 cv.cvSize (win_size, win_size), cv.cvSize (-1, -1),
 210                 cv.cvTermCriteria (cv.CV_TERMCRIT_ITER | cv.CV_TERMCRIT_EPS,
 211                                    20, 0.03))[0]
 212
 213             # we are no more in "add_remove_pt" mode
 214             add_remove_pt = False
 215
 216         # swapping
 217         prev_grey, grey = grey, prev_grey
 218         prev_pyramid, pyramid = pyramid, prev_pyramid
 219         points [0], points [1] = points [1], points [0]
 220         need_to_init = False
 221
 222         # we can now display the image
 223         highgui.cvShowImage ('LkDemo', image)
 224
 225         # handle events
 226         c = highgui.cvWaitKey (10)
 227
 228         if c == '\x1b':
 229             # user has press the ESC key, so exit
 230             break
 231
 232         # processing depending on the character
 233         if c in ['r', 'R']:
 234             need_to_init = True
 235         elif c in ['c', 'C']:
 236             points = [[], []]
 237         elif c in ['n', 'N']:
 238             night_mode = not night_mode