grabcut.py 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183
  1. #!/usr/bin/env python
  2. '''
  3. ===============================================================================
  4. Interactive Image Segmentation using GrabCut algorithm.
  5. This sample shows interactive image segmentation using grabcut algorithm.
  6. USAGE:
  7. python grabcut.py <filename>
  8. README FIRST:
  9. Two windows will show up, one for input and one for output.
  10. At first, in input window, draw a rectangle around the object using the
  11. right mouse button. Then press 'n' to segment the object (once or a few times)
  12. For any finer touch-ups, you can press any of the keys below and draw lines on
  13. the areas you want. Then again press 'n' to update the output.
  14. Key '0' - To select areas of sure background
  15. Key '1' - To select areas of sure foreground
  16. Key '2' - To select areas of probable background
  17. Key '3' - To select areas of probable foreground
  18. Key 'n' - To update the segmentation
  19. Key 'r' - To reset the setup
  20. Key 's' - To save the results
  21. ===============================================================================
  22. '''
  23. # Python 2/3 compatibility
  24. from __future__ import print_function
  25. import numpy as np
  26. import cv2 as cv
  27. import sys
  28. class App():
  29. BLUE = [255,0,0] # rectangle color
  30. RED = [0,0,255] # PR BG
  31. GREEN = [0,255,0] # PR FG
  32. BLACK = [0,0,0] # sure BG
  33. WHITE = [255,255,255] # sure FG
  34. DRAW_BG = {'color' : BLACK, 'val' : 0}
  35. DRAW_FG = {'color' : WHITE, 'val' : 1}
  36. DRAW_PR_BG = {'color' : RED, 'val' : 2}
  37. DRAW_PR_FG = {'color' : GREEN, 'val' : 3}
  38. # setting up flags
  39. rect = (0,0,1,1)
  40. drawing = False # flag for drawing curves
  41. rectangle = False # flag for drawing rect
  42. rect_over = False # flag to check if rect drawn
  43. rect_or_mask = 100 # flag for selecting rect or mask mode
  44. value = DRAW_FG # drawing initialized to FG
  45. thickness = 3 # brush thickness
  46. def onmouse(self, event, x, y, flags, param):
  47. # Draw Rectangle
  48. if event == cv.EVENT_RBUTTONDOWN:
  49. self.rectangle = True
  50. self.ix, self.iy = x,y
  51. elif event == cv.EVENT_MOUSEMOVE:
  52. if self.rectangle == True:
  53. self.img = self.img2.copy()
  54. cv.rectangle(self.img, (self.ix, self.iy), (x, y), self.BLUE, 2)
  55. self.rect = (min(self.ix, x), min(self.iy, y), abs(self.ix - x), abs(self.iy - y))
  56. self.rect_or_mask = 0
  57. elif event == cv.EVENT_RBUTTONUP:
  58. self.rectangle = False
  59. self.rect_over = True
  60. cv.rectangle(self.img, (self.ix, self.iy), (x, y), self.BLUE, 2)
  61. self.rect = (min(self.ix, x), min(self.iy, y), abs(self.ix - x), abs(self.iy - y))
  62. self.rect_or_mask = 0
  63. print(" Now press the key 'n' a few times until no further change \n")
  64. # draw touchup curves
  65. if event == cv.EVENT_LBUTTONDOWN:
  66. if self.rect_over == False:
  67. print("first draw rectangle \n")
  68. else:
  69. self.drawing = True
  70. cv.circle(self.img, (x,y), self.thickness, self.value['color'], -1)
  71. cv.circle(self.mask, (x,y), self.thickness, self.value['val'], -1)
  72. elif event == cv.EVENT_MOUSEMOVE:
  73. if self.drawing == True:
  74. cv.circle(self.img, (x, y), self.thickness, self.value['color'], -1)
  75. cv.circle(self.mask, (x, y), self.thickness, self.value['val'], -1)
  76. elif event == cv.EVENT_LBUTTONUP:
  77. if self.drawing == True:
  78. self.drawing = False
  79. cv.circle(self.img, (x, y), self.thickness, self.value['color'], -1)
  80. cv.circle(self.mask, (x, y), self.thickness, self.value['val'], -1)
  81. def run(self):
  82. # Loading images
  83. if len(sys.argv) == 2:
  84. filename = sys.argv[1] # for drawing purposes
  85. else:
  86. print("No input image given, so loading default image, lena.jpg \n")
  87. print("Correct Usage: python grabcut.py <filename> \n")
  88. filename = 'lena.jpg'
  89. self.img = cv.imread(cv.samples.findFile(filename))
  90. self.img2 = self.img.copy() # a copy of original image
  91. self.mask = np.zeros(self.img.shape[:2], dtype = np.uint8) # mask initialized to PR_BG
  92. self.output = np.zeros(self.img.shape, np.uint8) # output image to be shown
  93. # input and output windows
  94. cv.namedWindow('output')
  95. cv.namedWindow('input')
  96. cv.setMouseCallback('input', self.onmouse)
  97. cv.moveWindow('input', self.img.shape[1]+10,90)
  98. print(" Instructions: \n")
  99. print(" Draw a rectangle around the object using right mouse button \n")
  100. while(1):
  101. cv.imshow('output', self.output)
  102. cv.imshow('input', self.img)
  103. k = cv.waitKey(1)
  104. # key bindings
  105. if k == 27: # esc to exit
  106. break
  107. elif k == ord('0'): # BG drawing
  108. print(" mark background regions with left mouse button \n")
  109. self.value = self.DRAW_BG
  110. elif k == ord('1'): # FG drawing
  111. print(" mark foreground regions with left mouse button \n")
  112. self.value = self.DRAW_FG
  113. elif k == ord('2'): # PR_BG drawing
  114. self.value = self.DRAW_PR_BG
  115. elif k == ord('3'): # PR_FG drawing
  116. self.value = self.DRAW_PR_FG
  117. elif k == ord('s'): # save image
  118. bar = np.zeros((self.img.shape[0], 5, 3), np.uint8)
  119. res = np.hstack((self.img2, bar, self.img, bar, self.output))
  120. cv.imwrite('grabcut_output.png', res)
  121. print(" Result saved as image \n")
  122. elif k == ord('r'): # reset everything
  123. print("resetting \n")
  124. self.rect = (0,0,1,1)
  125. self.drawing = False
  126. self.rectangle = False
  127. self.rect_or_mask = 100
  128. self.rect_over = False
  129. self.value = self.DRAW_FG
  130. self.img = self.img2.copy()
  131. self.mask = np.zeros(self.img.shape[:2], dtype = np.uint8) # mask initialized to PR_BG
  132. self.output = np.zeros(self.img.shape, np.uint8) # output image to be shown
  133. elif k == ord('n'): # segment the image
  134. print(""" For finer touchups, mark foreground and background after pressing keys 0-3
  135. and again press 'n' \n""")
  136. try:
  137. bgdmodel = np.zeros((1, 65), np.float64)
  138. fgdmodel = np.zeros((1, 65), np.float64)
  139. if (self.rect_or_mask == 0): # grabcut with rect
  140. cv.grabCut(self.img2, self.mask, self.rect, bgdmodel, fgdmodel, 1, cv.GC_INIT_WITH_RECT)
  141. self.rect_or_mask = 1
  142. elif (self.rect_or_mask == 1): # grabcut with mask
  143. cv.grabCut(self.img2, self.mask, self.rect, bgdmodel, fgdmodel, 1, cv.GC_INIT_WITH_MASK)
  144. except:
  145. import traceback
  146. traceback.print_exc()
  147. mask2 = np.where((self.mask==1) + (self.mask==3), 255, 0).astype('uint8')
  148. self.output = cv.bitwise_and(self.img2, self.img2, mask=mask2)
  149. print('Done')
  150. if __name__ == '__main__':
  151. print(__doc__)
  152. App().run()
  153. cv.destroyAllWindows()