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