圖像分割：分水嶺分割算法--python實現

直接上代碼，原理大家都可以查到

注意：以下img自己加載即可，我是事先把圖片轉換成了單通道的Array形式，並且我的img是歸一化後在0-1的。

import cv2
import numpy as np
import copy
def watershed(img):

    gray = img
    ret0, thresh0 = cv2.threshold(gray,0,255,cv2.THRESH_BINARY_INV+cv2.THRESH_OTSU)

    kernel = np.ones((3,3),np.uint8)
    opening = cv2.morphologyEx(thresh0,cv2.MORPH_OPEN,kernel, iterations = 2)

    # 確定背景區域
    sure_bg = cv2.dilate(opening,kernel,iterations=3)

    # 確定前景區域
    dist_transform = cv2.distanceTransform(opening,cv2.DIST_L2,5)
    ret1, sure_fg = cv2.threshold(dist_transform,0.7*dist_transform.max(),255,0)

    # 查找未知區域
    sure_fg = np.uint8(sure_fg)
    unknown = cv2.subtract(sure_bg,sure_fg)

    # 標記標籤
    ret2, markers1 = cv2.connectedComponents(sure_fg)
    markers = markers1+1
    markers[unknown==255] = 0

    markers3 = cv2.watershed(img,markers)
    img[markers3 == -1] = [0,255,0]
    return thresh0,sure_bg,sure_fg,img

imgc = copy.deepcopy(img)
gray = copy.deepcopy(img)
ret, thresh = cv.threshold(gray, 0, 255, cv.THRESH_BINARY_INV + cv.THRESH_OTSU)
# noise removal
kernel = np.ones((3, 3), np.uint8)
opening = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=2)

# sure background area
sure_bg = cv2.dilate(opening, kernel, iterations=3)

# Finding sure foreground area
# 距離變換的基本含義是計算一個圖像中非零像素點到最近的零像素點的距離，也就是到零像素點的最短距離
# 最常見的距離變換算法就是通過連續的腐蝕操作來實現，腐蝕操作的停止條件是所有前景像素都被完全
# 腐蝕。這樣根據腐蝕的先後順序，我們就得到各個前景像素點到前景中心骨架像素點的距離。
# 根據各個像素點的距離值，設置爲不同的灰度值。這樣就完成了二值圖像的距離變換
# cv2.distanceTransform(src, distanceType, maskSize)
# 第二個參數 0,1,2 分別表示 CV_DIST_L1, CV_DIST_L2 , CV_DIST_C
dist_transform = cv2.distanceTransform(opening, 1, 5)
ret, sure_fg = cv2.threshold(dist_transform, 0.7 * dist_transform.max(), 255, 0)

# Finding unknown region
sure_fg = np.uint8(sure_fg)
unknown = cv2.subtract(sure_bg, sure_fg)
# Marker labelling
ret, markers = cv.connectedComponents(sure_fg)

# 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

imgc = cv.cvtColor(img, cv.COLOR_GRAY2BGR)
markers = cv.watershed(imgc, markers)

imgc[markers == -1] = [255,0,0]
label = imgc