Blame view

3rdparty/opencv-4.5.4/samples/python/lk_track.py 3.08 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
  #!/usr/bin/env python
  
  '''
  Lucas-Kanade tracker
  ====================
  
  Lucas-Kanade sparse optical flow demo. Uses goodFeaturesToTrack
  for track initialization and back-tracking for match verification
  between frames.
  
  Usage
  -----
  lk_track.py [<video_source>]
  
  
  Keys
  ----
  ESC - exit
  '''
  
  # Python 2/3 compatibility
  from __future__ import print_function
  
  import numpy as np
  import cv2 as cv
  
  import video
  from common import anorm2, draw_str
  
  lk_params = dict( winSize  = (15, 15),
                    maxLevel = 2,
                    criteria = (cv.TERM_CRITERIA_EPS | cv.TERM_CRITERIA_COUNT, 10, 0.03))
  
  feature_params = dict( maxCorners = 500,
                         qualityLevel = 0.3,
                         minDistance = 7,
                         blockSize = 7 )
  
  class App:
      def __init__(self, video_src):
          self.track_len = 10
          self.detect_interval = 5
          self.tracks = []
          self.cam = video.create_capture(video_src)
          self.frame_idx = 0
  
      def run(self):
          while True:
              _ret, frame = self.cam.read()
              frame_gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
              vis = frame.copy()
  
              if len(self.tracks) > 0:
                  img0, img1 = self.prev_gray, frame_gray
                  p0 = np.float32([tr[-1] for tr in self.tracks]).reshape(-1, 1, 2)
                  p1, _st, _err = cv.calcOpticalFlowPyrLK(img0, img1, p0, None, **lk_params)
                  p0r, _st, _err = cv.calcOpticalFlowPyrLK(img1, img0, p1, None, **lk_params)
                  d = abs(p0-p0r).reshape(-1, 2).max(-1)
                  good = d < 1
                  new_tracks = []
                  for tr, (x, y), good_flag in zip(self.tracks, p1.reshape(-1, 2), good):
                      if not good_flag:
                          continue
                      tr.append((x, y))
                      if len(tr) > self.track_len:
                          del tr[0]
                      new_tracks.append(tr)
                      cv.circle(vis, (int(x), int(y)), 2, (0, 255, 0), -1)
                  self.tracks = new_tracks
                  cv.polylines(vis, [np.int32(tr) for tr in self.tracks], False, (0, 255, 0))
                  draw_str(vis, (20, 20), 'track count: %d' % len(self.tracks))
  
              if self.frame_idx % self.detect_interval == 0:
                  mask = np.zeros_like(frame_gray)
                  mask[:] = 255
                  for x, y in [np.int32(tr[-1]) for tr in self.tracks]:
                      cv.circle(mask, (x, y), 5, 0, -1)
                  p = cv.goodFeaturesToTrack(frame_gray, mask = mask, **feature_params)
                  if p is not None:
                      for x, y in np.float32(p).reshape(-1, 2):
                          self.tracks.append([(x, y)])
  
  
              self.frame_idx += 1
              self.prev_gray = frame_gray
              cv.imshow('lk_track', vis)
  
              ch = cv.waitKey(1)
              if ch == 27:
                  break
  
  def main():
      import sys
      try:
          video_src = sys.argv[1]
      except:
          video_src = 0
  
      App(video_src).run()
      print('Done')
  
  
  if __name__ == '__main__':
      print(__doc__)
      main()
      cv.destroyAllWindows()