Home | History | Annotate | Download | only in python2
      1 #!/usr/bin/env python
      2 
      3 '''
      4 Multitarget planar tracking
      5 ==================
      6 
      7 Example of using features2d framework for interactive video homography matching.
      8 ORB features and FLANN matcher are used. This sample provides PlaneTracker class
      9 and an example of its usage.
     10 
     11 video: http://www.youtube.com/watch?v=pzVbhxx6aog
     12 
     13 Usage
     14 -----
     15 plane_tracker.py [<video source>]
     16 
     17 Keys:
     18    SPACE  -  pause video
     19    c      -  clear targets
     20 
     21 Select a textured planar object to track by drawing a box with a mouse.
     22 '''
     23 
     24 import numpy as np
     25 import cv2
     26 
     27 # built-in modules
     28 from collections import namedtuple
     29 
     30 # local modules
     31 import video
     32 import common
     33 
     34 
     35 FLANN_INDEX_KDTREE = 1
     36 FLANN_INDEX_LSH    = 6
     37 flann_params= dict(algorithm = FLANN_INDEX_LSH,
     38                    table_number = 6, # 12
     39                    key_size = 12,     # 20
     40                    multi_probe_level = 1) #2
     41 
     42 MIN_MATCH_COUNT = 10
     43 
     44 '''
     45   image     - image to track
     46   rect      - tracked rectangle (x1, y1, x2, y2)
     47   keypoints - keypoints detected inside rect
     48   descrs    - their descriptors
     49   data      - some user-provided data
     50 '''
     51 PlanarTarget = namedtuple('PlaneTarget', 'image, rect, keypoints, descrs, data')
     52 
     53 '''
     54   target - reference to PlanarTarget
     55   p0     - matched points coords in target image
     56   p1     - matched points coords in input frame
     57   H      - homography matrix from p0 to p1
     58   quad   - target bounary quad in input frame
     59 '''
     60 TrackedTarget = namedtuple('TrackedTarget', 'target, p0, p1, H, quad')
     61 
     62 class PlaneTracker:
     63     def __init__(self):
     64         self.detector = cv2.ORB_create( nfeatures = 1000 )
     65         self.matcher = cv2.FlannBasedMatcher(flann_params, {})  # bug : need to pass empty dict (#1329)
     66         self.targets = []
     67 
     68     def add_target(self, image, rect, data=None):
     69         '''Add a new tracking target.'''
     70         x0, y0, x1, y1 = rect
     71         raw_points, raw_descrs = self.detect_features(image)
     72         points, descs = [], []
     73         for kp, desc in zip(raw_points, raw_descrs):
     74             x, y = kp.pt
     75             if x0 <= x <= x1 and y0 <= y <= y1:
     76                 points.append(kp)
     77                 descs.append(desc)
     78         descs = np.uint8(descs)
     79         self.matcher.add([descs])
     80         target = PlanarTarget(image = image, rect=rect, keypoints = points, descrs=descs, data=data)
     81         self.targets.append(target)
     82 
     83     def clear(self):
     84         '''Remove all targets'''
     85         self.targets = []
     86         self.matcher.clear()
     87 
     88     def track(self, frame):
     89         '''Returns a list of detected TrackedTarget objects'''
     90         frame_points, frame_descrs = self.detect_features(frame)
     91         if len(frame_points) < MIN_MATCH_COUNT:
     92             return []
     93         matches = self.matcher.knnMatch(frame_descrs, k = 2)
     94         matches = [m[0] for m in matches if len(m) == 2 and m[0].distance < m[1].distance * 0.75]
     95         if len(matches) < MIN_MATCH_COUNT:
     96             return []
     97         matches_by_id = [[] for _ in xrange(len(self.targets))]
     98         for m in matches:
     99             matches_by_id[m.imgIdx].append(m)
    100         tracked = []
    101         for imgIdx, matches in enumerate(matches_by_id):
    102             if len(matches) < MIN_MATCH_COUNT:
    103                 continue
    104             target = self.targets[imgIdx]
    105             p0 = [target.keypoints[m.trainIdx].pt for m in matches]
    106             p1 = [frame_points[m.queryIdx].pt for m in matches]
    107             p0, p1 = np.float32((p0, p1))
    108             H, status = cv2.findHomography(p0, p1, cv2.RANSAC, 3.0)
    109             status = status.ravel() != 0
    110             if status.sum() < MIN_MATCH_COUNT:
    111                 continue
    112             p0, p1 = p0[status], p1[status]
    113 
    114             x0, y0, x1, y1 = target.rect
    115             quad = np.float32([[x0, y0], [x1, y0], [x1, y1], [x0, y1]])
    116             quad = cv2.perspectiveTransform(quad.reshape(1, -1, 2), H).reshape(-1, 2)
    117 
    118             track = TrackedTarget(target=target, p0=p0, p1=p1, H=H, quad=quad)
    119             tracked.append(track)
    120         tracked.sort(key = lambda t: len(t.p0), reverse=True)
    121         return tracked
    122 
    123     def detect_features(self, frame):
    124         '''detect_features(self, frame) -> keypoints, descrs'''
    125         keypoints, descrs = self.detector.detectAndCompute(frame, None)
    126         if descrs is None:  # detectAndCompute returns descs=None if not keypoints found
    127             descrs = []
    128         return keypoints, descrs
    129 
    130 
    131 class App:
    132     def __init__(self, src):
    133         self.cap = video.create_capture(src)
    134         self.frame = None
    135         self.paused = False
    136         self.tracker = PlaneTracker()
    137 
    138         cv2.namedWindow('plane')
    139         self.rect_sel = common.RectSelector('plane', self.on_rect)
    140 
    141     def on_rect(self, rect):
    142         self.tracker.add_target(self.frame, rect)
    143 
    144     def run(self):
    145         while True:
    146             playing = not self.paused and not self.rect_sel.dragging
    147             if playing or self.frame is None:
    148                 ret, frame = self.cap.read()
    149                 if not ret:
    150                     break
    151                 self.frame = frame.copy()
    152 
    153             vis = self.frame.copy()
    154             if playing:
    155                 tracked = self.tracker.track(self.frame)
    156                 for tr in tracked:
    157                     cv2.polylines(vis, [np.int32(tr.quad)], True, (255, 255, 255), 2)
    158                     for (x, y) in np.int32(tr.p1):
    159                         cv2.circle(vis, (x, y), 2, (255, 255, 255))
    160 
    161             self.rect_sel.draw(vis)
    162             cv2.imshow('plane', vis)
    163             ch = cv2.waitKey(1) & 0xFF
    164             if ch == ord(' '):
    165                 self.paused = not self.paused
    166             if ch == ord('c'):
    167                 self.tracker.clear()
    168             if ch == 27:
    169                 break
    170 
    171 if __name__ == '__main__':
    172     print __doc__
    173 
    174     import sys
    175     try:
    176         video_src = sys.argv[1]
    177     except:
    178         video_src = 0
    179     App(video_src).run()
    180