欢迎访问宙启技术站
智能推送

使用Python实现watershed(分水岭)算法进行图像分割的代码示例

发布时间:2023-12-27 13:28:51

Watershed算法是一种基于图论的图像分割算法,可以将图像分割成具有明显边界的不同区域。下面是使用Python实现Watershed算法进行图像分割的代码示例:

import cv2
import numpy as np

def watershed_segmentation(image_path):
    # 读取图像
    image = cv2.imread(image_path)
    
    # 将图像转换为灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    # 阈值分割得到二值图像
    ret, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV+cv2.THRESH_OTSU)
    
    # 对二值图像进行形态学操作,去除噪声点
    kernel = np.ones((3,3), np.uint8)
    opening = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=2)
    
    # 找到背景区域
    sure_bg = cv2.dilate(opening, kernel, iterations=3)
    
    # 距离变换得到前景区域
    dist_transform = cv2.distanceTransform(opening, cv2.DIST_L2, 5)
    ret, 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)
    
    # 标记未知区域
    ret, markers = cv2.connectedComponents(sure_fg)
    markers = markers + 1
    markers[unknown == 255] = 0
    
    # 使用分水岭算法进行图像分割
    markers = cv2.watershed(image, markers)
    image[markers == -1] = [0, 0, 255]
    
    return image

# 使用示例
image_path = 'example.jpg'
segmented_image = watershed_segmentation(image_path)

# 显示原始图像和分割结果
cv2.imshow('Original Image', cv2.imread(image_path))
cv2.imshow('Segmented Image', segmented_image)
cv2.waitKey(0)
cv2.destroyAllWindows()

上述代码中,我们首先读取图像,并将其转换为灰度图。然后使用阈值分割得到二值图像,并对二值图像进行形态学操作来去除噪声。接下来,使用距离变换得到前景区域,并找到未知区域。然后,我们对未知区域进行标记,并使用分水岭算法进行图像分割。最后,将分割结果显示出来。

需要注意的是,Watershed算法需要提前确定一些标记,例如确定背景区域和前景区域。在代码示例中,我们使用阈值分割和形态学操作来获取这些标记。对于不同的图像,可能需要调整阈值等参数以获得更好的分割效果。