samples/python/watershed.py

   1 #!/usr/bin/env python
   2
   3 '''
   4 Watershed segmentation
   5 =========
   6
   7 This program demonstrates the watershed segmentation algorithm
   8 in OpenCV: watershed().
   9
  10 Usage
  11 -----
  12 watershed.py [image filename]
  13
  14 Keys
  15 ----
  16   1-7   - switch marker color
  17   SPACE - update segmentation
  18   r     - reset
  19   a     - toggle autoupdate
  20   ESC   - exit
  21
  22 '''
  23
  24
  25 # Python 2/3 compatibility
  26 from __future__ import print_function
  27
  28 import numpy as np
  29 import cv2 as cv
  30 from common import Sketcher
  31
  32 class App:
  33     def __init__(self, fn):
  34         self.img = cv.imread(fn)
  35         if self.img is None:
  36             raise Exception('Failed to load image file: %s' % fn)
  37
  38         h, w = self.img.shape[:2]
  39         self.markers = np.zeros((h, w), np.int32)
  40         self.markers_vis = self.img.copy()
  41         self.cur_marker = 1
  42         self.colors = np.int32( list(np.ndindex(2, 2, 2)) ) * 255
  43
  44         self.auto_update = True
  45         self.sketch = Sketcher('img', [self.markers_vis, self.markers], self.get_colors)
  46
  47     def get_colors(self):
  48         return list(map(int, self.colors[self.cur_marker])), self.cur_marker
  49
  50     def watershed(self):
  51         m = self.markers.copy()
  52         cv.watershed(self.img, m)
  53         overlay = self.colors[np.maximum(m, 0)]
  54         vis = cv.addWeighted(self.img, 0.5, overlay, 0.5, 0.0, dtype=cv.CV_8UC3)
  55         cv.imshow('watershed', vis)
  56
  57     def run(self):
  58         while cv.getWindowProperty('img', 0) != -1 or cv.getWindowProperty('watershed', 0) != -1:
  59             ch = cv.waitKey(50)
  60             if ch == 27:
  61                 break
  62             if ch >= ord('1') and ch <= ord('7'):
  63                 self.cur_marker = ch - ord('0')
  64                 print('marker: ', self.cur_marker)
  65             if ch == ord(' ') or (self.sketch.dirty and self.auto_update):
  66                 self.watershed()
  67                 self.sketch.dirty = False
  68             if ch in [ord('a'), ord('A')]:
  69                 self.auto_update = not self.auto_update
  70                 print('auto_update if', ['off', 'on'][self.auto_update])
  71             if ch in [ord('r'), ord('R')]:
  72                 self.markers[:] = 0
  73                 self.markers_vis[:] = self.img
  74                 self.sketch.show()
  75         cv.destroyAllWindows()
  76
  77
  78 if __name__ == '__main__':
  79     import sys
  80     try:
  81         fn = sys.argv[1]
  82     except:
  83         fn = 'fruits.jpg'
  84     print(__doc__)
  85     App(cv.samples.findFile(fn)).run()