-
Notifications
You must be signed in to change notification settings - Fork 0
/
watershedOpenCV.py
53 lines (39 loc) · 1.55 KB
/
watershedOpenCV.py
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
import cv2 as cv
import numpy as np
def watershedSegmentation(img):
"""Segment the image by distinguishing the foreground and the background
regions, and then use them to compute a segmentation using watershed.
Args:
img (numpy array): input image.
Returns:
numpy array: segmented image.
"""
# Convert to gray from bgr
grayImg = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
# Apply a Gaussian Blur to the grayImg to remove noise
grayImg = cv.GaussianBlur(grayImg, (3, 3), 0)
# Threshold the gray image
thresh = cv.threshold(
grayImg,
0,
255,
cv.THRESH_BINARY_INV + cv.THRESH_OTSU,
)[1]
# Finding sure foreground area
distTransform = cv.distanceTransform(thresh, cv.DIST_L2, 5)
sureForeground = cv.threshold(distTransform, 0.65 * distTransform.max(), 255, 0)[1]
sureForeground = np.uint8(sureForeground)
# Finding unknown region
unknown = cv.subtract(thresh, sureForeground)
# Marker labelling
markers = cv.connectedComponents(sureForeground)[1]
# Add one to all labels so that sure background is not 0, but 1
markers = markers + 1
# Now, mark the region of unknown with zero
markers[unknown == 255] = 0
# Apply watershed to modify the markers, fiding the final regions.
markers = cv.watershed(img, markers)
# Use the found regions and mark them with black and white on the input image.
img[markers > 1] = [0, 0, 0]
img[markers <= 1] = [255, 255, 255]
return img