samples/python/snippets/watershed.py#

An example using the watershed algorithm using python

 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# Python 2/3 compatibility
26from __future__ import print_function
27
28import numpy as np
29import cv2 as cv
30from common import Sketcher
31
32class 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
78if __name__ == '__main__':
79    print(__doc__)
80    import sys
81    try:
82        fn = sys.argv[1]
83    except:
84        fn = 'fruits.jpg'
85    App(cv.samples.findFile(fn)).run()