Blame view

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