# 分水岭算法分割图像
# 1、图像预处理：对原始图像进行预处理，例如将彩色图像转换为灰度图像，平滑图像以去除噪声，增强图像对比度等。
# 2、计算图像的梯度：通过对图像应用梯度算子（如Sobel算子）来计算图像的梯度，梯度的变化反映了图像中的边缘信息。
# 3、确定种子点：根据预处理后的图像或者用户定义的标记信息，确定一些种子点。这些种子点可以代表分割的目标物体或者背景。
# 4、计算距离变换：使用距离变换算法，从种子点开始，计算每个像素点到最近种子点的距离。这样可以得到一个距离图像，表示每个像素点离最近种子点的距离。
# 5、计算梯度和距离变换的合并图像：将梯度图像和距离变换图像进行合并，得到一个新的图像。在合并图像中，水流将被认为是从梯度较高的区域向梯度较低的区域流去。
# 6、计算分水岭：根据合并图像，应用分水岭变换算法来计算分水岭，即将图像分割成不同的分水岭区域。
# 7、降噪和后处理：对分割结果进行降噪和后处理，例如去除小的分割区域、填充空洞等，以获得最终的分割结果。

import cv2 as cv
import numpy as np

img = cv.imread('test1.png', cv.IMREAD_GRAYSCALE)
img = cv.blur(img, (5, 5))
img = cv.convertScaleAbs(img, alpha=1.0, beta=2.0)
cv.imshow('origin', img)

grad_x = cv.Sobel(img, cv.CV_32F, 1, 0, ksize=3)
grad_y = cv.Sobel(img, cv.CV_32F, 0, 1, ksize=3)
grad = cv.magnitude(grad_x, grad_y)

ret, markers = cv.threshold(grad, 50, 255, cv.THRESH_BINARY)

markers = np.uint8(markers)
dist_transform = cv.distanceTransform(255 - markers, cv.DIST_L2, 5)

combined = cv.subtract(dist_transform, grad)

image_rgb = cv.cvtColor(cv.imread('test1.png'), cv.COLOR_BGR2RGB)
markers = np.int32(markers)
markers[grad == 0] = 0
markers += 1
cv.watershed(image_rgb, markers)

markers = np.where(markers == -1, 255, markers)
markers = np.uint8(markers)

output = np.zeros_like(image_rgb)
output[markers == 255] = [0, 0, 255]
output[markers == 1] = [18, 0, 230]
output[markers == 0] = [255, 255, 255]

cv.imshow('Segmentation Result', output)
cv.waitKey(0)
cv.destroyAllWindows()