Blame view

3rdparty/opencv-4.5.4/samples/python/grabcut.py 7.33 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
  #!/usr/bin/env python
  '''
  ===============================================================================
  Interactive Image Segmentation using GrabCut algorithm.
  
  This sample shows interactive image segmentation using grabcut algorithm.
  
  USAGE:
      python grabcut.py <filename>
  
  README FIRST:
      Two windows will show up, one for input and one for output.
  
      At first, in input window, draw a rectangle around the object using the
  right mouse button. Then press 'n' to segment the object (once or a few times)
  For any finer touch-ups, you can press any of the keys below and draw lines on
  the areas you want. Then again press 'n' to update the output.
  
  Key '0' - To select areas of sure background
  Key '1' - To select areas of sure foreground
  Key '2' - To select areas of probable background
  Key '3' - To select areas of probable foreground
  
  Key 'n' - To update the segmentation
  Key 'r' - To reset the setup
  Key 's' - To save the results
  ===============================================================================
  '''
  
  # Python 2/3 compatibility
  from __future__ import print_function
  
  import numpy as np
  import cv2 as cv
  
  import sys
  
  class App():
      BLUE = [255,0,0]        # rectangle color
      RED = [0,0,255]         # PR BG
      GREEN = [0,255,0]       # PR FG
      BLACK = [0,0,0]         # sure BG
      WHITE = [255,255,255]   # sure FG
  
      DRAW_BG = {'color' : BLACK, 'val' : 0}
      DRAW_FG = {'color' : WHITE, 'val' : 1}
      DRAW_PR_BG = {'color' : RED, 'val' : 2}
      DRAW_PR_FG = {'color' : GREEN, 'val' : 3}
  
      # setting up flags
      rect = (0,0,1,1)
      drawing = False         # flag for drawing curves
      rectangle = False       # flag for drawing rect
      rect_over = False       # flag to check if rect drawn
      rect_or_mask = 100      # flag for selecting rect or mask mode
      value = DRAW_FG         # drawing initialized to FG
      thickness = 3           # brush thickness
  
      def onmouse(self, event, x, y, flags, param):
          # Draw Rectangle
          if event == cv.EVENT_RBUTTONDOWN:
              self.rectangle = True
              self.ix, self.iy = x,y
  
          elif event == cv.EVENT_MOUSEMOVE:
              if self.rectangle == True:
                  self.img = self.img2.copy()
                  cv.rectangle(self.img, (self.ix, self.iy), (x, y), self.BLUE, 2)
                  self.rect = (min(self.ix, x), min(self.iy, y), abs(self.ix - x), abs(self.iy - y))
                  self.rect_or_mask = 0
  
          elif event == cv.EVENT_RBUTTONUP:
              self.rectangle = False
              self.rect_over = True
              cv.rectangle(self.img, (self.ix, self.iy), (x, y), self.BLUE, 2)
              self.rect = (min(self.ix, x), min(self.iy, y), abs(self.ix - x), abs(self.iy - y))
              self.rect_or_mask = 0
              print(" Now press the key 'n' a few times until no further change \n")
  
          # draw touchup curves
  
          if event == cv.EVENT_LBUTTONDOWN:
              if self.rect_over == False:
                  print("first draw rectangle \n")
              else:
                  self.drawing = True
                  cv.circle(self.img, (x,y), self.thickness, self.value['color'], -1)
                  cv.circle(self.mask, (x,y), self.thickness, self.value['val'], -1)
  
          elif event == cv.EVENT_MOUSEMOVE:
              if self.drawing == True:
                  cv.circle(self.img, (x, y), self.thickness, self.value['color'], -1)
                  cv.circle(self.mask, (x, y), self.thickness, self.value['val'], -1)
  
          elif event == cv.EVENT_LBUTTONUP:
              if self.drawing == True:
                  self.drawing = False
                  cv.circle(self.img, (x, y), self.thickness, self.value['color'], -1)
                  cv.circle(self.mask, (x, y), self.thickness, self.value['val'], -1)
  
      def run(self):
          # Loading images
          if len(sys.argv) == 2:
              filename = sys.argv[1] # for drawing purposes
          else:
              print("No input image given, so loading default image, lena.jpg \n")
              print("Correct Usage: python grabcut.py <filename> \n")
              filename = 'lena.jpg'
  
          self.img = cv.imread(cv.samples.findFile(filename))
          self.img2 = self.img.copy()                               # a copy of original image
          self.mask = np.zeros(self.img.shape[:2], dtype = np.uint8) # mask initialized to PR_BG
          self.output = np.zeros(self.img.shape, np.uint8)           # output image to be shown
  
          # input and output windows
          cv.namedWindow('output')
          cv.namedWindow('input')
          cv.setMouseCallback('input', self.onmouse)
          cv.moveWindow('input', self.img.shape[1]+10,90)
  
          print(" Instructions: \n")
          print(" Draw a rectangle around the object using right mouse button \n")
  
          while(1):
  
              cv.imshow('output', self.output)
              cv.imshow('input', self.img)
              k = cv.waitKey(1)
  
              # key bindings
              if k == 27:         # esc to exit
                  break
              elif k == ord('0'): # BG drawing
                  print(" mark background regions with left mouse button \n")
                  self.value = self.DRAW_BG
              elif k == ord('1'): # FG drawing
                  print(" mark foreground regions with left mouse button \n")
                  self.value = self.DRAW_FG
              elif k == ord('2'): # PR_BG drawing
                  self.value = self.DRAW_PR_BG
              elif k == ord('3'): # PR_FG drawing
                  self.value = self.DRAW_PR_FG
              elif k == ord('s'): # save image
                  bar = np.zeros((self.img.shape[0], 5, 3), np.uint8)
                  res = np.hstack((self.img2, bar, self.img, bar, self.output))
                  cv.imwrite('grabcut_output.png', res)
                  print(" Result saved as image \n")
              elif k == ord('r'): # reset everything
                  print("resetting \n")
                  self.rect = (0,0,1,1)
                  self.drawing = False
                  self.rectangle = False
                  self.rect_or_mask = 100
                  self.rect_over = False
                  self.value = self.DRAW_FG
                  self.img = self.img2.copy()
                  self.mask = np.zeros(self.img.shape[:2], dtype = np.uint8) # mask initialized to PR_BG
                  self.output = np.zeros(self.img.shape, np.uint8)           # output image to be shown
              elif k == ord('n'): # segment the image
                  print(""" For finer touchups, mark foreground and background after pressing keys 0-3
                  and again press 'n' \n""")
                  try:
                      bgdmodel = np.zeros((1, 65), np.float64)
                      fgdmodel = np.zeros((1, 65), np.float64)
                      if (self.rect_or_mask == 0):         # grabcut with rect
                          cv.grabCut(self.img2, self.mask, self.rect, bgdmodel, fgdmodel, 1, cv.GC_INIT_WITH_RECT)
                          self.rect_or_mask = 1
                      elif (self.rect_or_mask == 1):       # grabcut with mask
                          cv.grabCut(self.img2, self.mask, self.rect, bgdmodel, fgdmodel, 1, cv.GC_INIT_WITH_MASK)
                  except:
                      import traceback
                      traceback.print_exc()
  
              mask2 = np.where((self.mask==1) + (self.mask==3), 255, 0).astype('uint8')
              self.output = cv.bitwise_and(self.img2, self.img2, mask=mask2)
  
          print('Done')
  
  
  if __name__ == '__main__':
      print(__doc__)
      App().run()
      cv.destroyAllWindows()