Blame view

3rdparty/opencv-4.5.4/samples/python/dis_opt_flow.py 3.46 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
  #!/usr/bin/env python
  
  '''
  example to show optical flow estimation using DISOpticalFlow
  
  USAGE: dis_opt_flow.py [<video_source>]
  
  Keys:
   1  - toggle HSV flow visualization
   2  - toggle glitch
   3  - toggle spatial propagation of flow vectors
   4  - toggle temporal propagation of flow vectors
  ESC - exit
  '''
  
  # Python 2/3 compatibility
  from __future__ import print_function
  
  import numpy as np
  import cv2 as cv
  
  import video
  
  
  def draw_flow(img, flow, step=16):
      h, w = img.shape[:2]
      y, x = np.mgrid[step/2:h:step, step/2:w:step].reshape(2,-1).astype(int)
      fx, fy = flow[y,x].T
      lines = np.vstack([x, y, x+fx, y+fy]).T.reshape(-1, 2, 2)
      lines = np.int32(lines + 0.5)
      vis = cv.cvtColor(img, cv.COLOR_GRAY2BGR)
      cv.polylines(vis, lines, 0, (0, 255, 0))
      for (x1, y1), (_x2, _y2) in lines:
          cv.circle(vis, (x1, y1), 1, (0, 255, 0), -1)
      return vis
  
  
  def draw_hsv(flow):
      h, w = flow.shape[:2]
      fx, fy = flow[:,:,0], flow[:,:,1]
      ang = np.arctan2(fy, fx) + np.pi
      v = np.sqrt(fx*fx+fy*fy)
      hsv = np.zeros((h, w, 3), np.uint8)
      hsv[...,0] = ang*(180/np.pi/2)
      hsv[...,1] = 255
      hsv[...,2] = np.minimum(v*4, 255)
      bgr = cv.cvtColor(hsv, cv.COLOR_HSV2BGR)
      return bgr
  
  
  def warp_flow(img, flow):
      h, w = flow.shape[:2]
      flow = -flow
      flow[:,:,0] += np.arange(w)
      flow[:,:,1] += np.arange(h)[:,np.newaxis]
      res = cv.remap(img, flow, None, cv.INTER_LINEAR)
      return res
  
  
  def main():
      import sys
      print(__doc__)
      try:
          fn = sys.argv[1]
      except IndexError:
          fn = 0
  
      cam = video.create_capture(fn)
      _ret, prev = cam.read()
      prevgray = cv.cvtColor(prev, cv.COLOR_BGR2GRAY)
      show_hsv = False
      show_glitch = False
      use_spatial_propagation = False
      use_temporal_propagation = True
      cur_glitch = prev.copy()
      inst = cv.DISOpticalFlow.create(cv.DISOPTICAL_FLOW_PRESET_MEDIUM)
      inst.setUseSpatialPropagation(use_spatial_propagation)
  
      flow = None
      while True:
          _ret, img = cam.read()
          gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
          if flow is not None and use_temporal_propagation:
              #warp previous flow to get an initial approximation for the current flow:
              flow = inst.calc(prevgray, gray, warp_flow(flow,flow))
          else:
              flow = inst.calc(prevgray, gray, None)
          prevgray = gray
  
          cv.imshow('flow', draw_flow(gray, flow))
          if show_hsv:
              cv.imshow('flow HSV', draw_hsv(flow))
          if show_glitch:
              cur_glitch = warp_flow(cur_glitch, flow)
              cv.imshow('glitch', cur_glitch)
  
          ch = 0xFF & cv.waitKey(5)
          if ch == 27:
              break
          if ch == ord('1'):
              show_hsv = not show_hsv
              print('HSV flow visualization is', ['off', 'on'][show_hsv])
          if ch == ord('2'):
              show_glitch = not show_glitch
              if show_glitch:
                  cur_glitch = img.copy()
              print('glitch is', ['off', 'on'][show_glitch])
          if ch == ord('3'):
              use_spatial_propagation = not use_spatial_propagation
              inst.setUseSpatialPropagation(use_spatial_propagation)
              print('spatial propagation is', ['off', 'on'][use_spatial_propagation])
          if ch == ord('4'):
              use_temporal_propagation = not use_temporal_propagation
              print('temporal propagation is', ['off', 'on'][use_temporal_propagation])
  
      print('Done')
  
  
  if __name__ == '__main__':
      print(__doc__)
      main()
      cv.destroyAllWindows()