python手撕分水岭算法

1 分水岭算法实现

主要思路就是：

利用一个优先队列与有序队列（有序队列其实可以不用）。优先队列是按像素的灰度值排列的，灰度值低的先被淹。
通过统计像素的附近的点的标记种类个数来确认当前像素点的标记，并把该像素的四邻域入队。
注意要使用梯度图像，这样使得各个部分蔓延的速度一样。

import cv2
import numpy as np
import matplotlib.pyplot as plt
import random
import plotly.express as px
from queue import Queue,PriorityQueue
import sys

class WATERSHED:
    def __init__(self,img,markers):
        if img.ndim!=2:
            raise ValueError("请输入灰度图")
        self.img = img
        # 求图像梯度幅值
        self.gradient()
        self.markers = markers.astype(np.int32)
        self.img_highest = np.max(self.img)
        self.img_lowest = np.min(self.img)
        self.WSHED = -1# 分水岭
        self.IN_QUEUE = -2#已入队，但未赋予标记
        # 根据梯度幅值图的灰度最大最小值创建有序队列
        self.ordered_queue = [[] for _ in range(self.img_highest-self.img_lowest+1)] 
    def gradient(self):
        # self.img = cv2.Canny(self.img,80,150)
        d_x= cv2.Sobel(self.img,cv2.CV_64F,1,0,ksize=3)
        d_y= cv2.Sobel(self.img,cv2.CV_64F,0,1,ksize=3)
        value = cv2.magnitude(d_x,d_y)
        self.img = cv2.convertScaleAbs(value)
    def ordered_queue_push(self,i,j):
        gray_leve = self.img[i][j]
        self.ordered_queue[gray_leve-self.img_lowest].append([i,j])
        self.markers[i][j] = self.IN_QUEUE
    def PriorityQueue_push(self,water_stage):
        # 将等于水位高度的像素点全部入优先队列
        for pix_coord in self.ordered_queue[water_stage-self.img_lowest]:
            self.wait_water_queue.put((water_stage,pix_coord))
    def wait_water_queue_push(self,x,y):
        self.wait_water_queue.put((self.img[x][y],[x,y]))
        self.markers[x][y] = self.IN_QUEUE
    def get_label(self,pix_coord:list,water_stage):
        x,y = pix_coord
        neighbour_pix_values = {
    
    }
        for i in [-1,0,1]:
            for j in [-1,0,1]:
                pix_x,pix_y = x+i,y+j
                if pix_x<0 or pix_y<0 or pix_x >=self.img.shape[0] \
                or pix_y >=self.img.shape[1] :
                    continue
                # 必须以四邻域搜索，八邻域会越过边界
                if abs(i+j)==1:
                    if neighbour_pix_values.get(self.markers[pix_x][pix_y])!=None:
                        neighbour_pix_values[self.markers[pix_x][pix_y]]+=1
                    else:
                        neighbour_pix_values[self.markers[pix_x][pix_y]]=1
                    if self.markers[pix_x][pix_y]==0:#不在队列且无标签
                        if self.img[pix_x][pix_y]<=water_stage:
                            #低于水位，入优先队列
                            self.wait_water_queue_push(pix_x,pix_y)
                        else:
                            #高于水位，还淹不到
                            self.ordered_queue_push(pix_x,pix_y)
        
        pix_values = list(neighbour_pix_values)
        #去掉分水岭、入队标记和未标记区域值
        pix_values = [value for value in pix_values \
            if value>0 and value!=self.IN_QUEUE]

        if len(pix_values)==1:           
            self.markers[x][y]=pix_values[0]
        else :
            self.markers[x][y]=self.WSHED            
    def run(self):
        
        for i in range(self.img.shape[0]):
            for j in range(self.img.shape[1]):
                #已标记区域全部入队
                if self.markers[i][j] >0:
                    gray_leve = self.img[i][j]
                    self.ordered_queue[gray_leve-self.img_lowest].append([i,j])
                           
        for water_stage in range(self.img_lowest,self.img_highest+1):
            sys.stdout.write("水位高度{}".format(water_stage))
            sys.stdout.flush()
            self.wait_water_queue = PriorityQueue()
            #将水位为water_stage的全部入优先队列
            self.PriorityQueue_push(water_stage)
            #每轮循环结束小于等于water_stage的像素值都得清空
            while self.wait_water_queue.qsize() > 0:
                pix_coord = self.wait_water_queue.get()[1]
                self.get_label(pix_coord,water_stage)
            # 动画
            plt.clf()
            plt.imshow(self.markers)
            plt.pause(0.001)
        return self.markers

2 获取初始标记

如果如opencv官方例程那样，用距离变换来确认标记，非常容易过分割，所以我们可以通过手动进行标记。
可以使用opencv的鼠标事件实现标记。

import cv2
import numpy as np

class draw_markers:
    def __init__(self,img,resize_rate:int=1):
        if img.ndim == 2:
            self.img = cv2.cvtColor(img,cv2.COLOR_GRAY2BGR)
        else:
            self.img = img
        self. resize_rate = resize_rate
        
        self.r,self.w = self.img.shape[:2]
        self.markers = np.zeros_like(self.img)
        # 放大
        self.img = cv2.resize(self.img,(self.r*resize_rate,self.w*resize_rate))
        self.markers = cv2.resize(self.markers,(self.r*resize_rate,self.w*resize_rate))
        
        self.cls_num =0
        #标记颜色列表
        self.color_list = [(255,0,0),(0,255,0),(0,0,255),(255,255,0),(255,0,255),(0,255,255)]
        self.drawing = False
    def mark(self,event, x, y,flags, param):
        if event==cv2.EVENT_LBUTTONDOWN:
            self.drawing = True
        if event == cv2.EVENT_LBUTTONUP:
            self.drawing = False
        if self.drawing:
            cv2.circle(self.img,(x,y),radius=1,color=self.color_list[self.cls_num],thickness=-1)
            cv2.circle(self.markers,(x,y),radius=1,color=self.color_list[self.cls_num],thickness=-1)
    def draw(self):
        cv2.namedWindow('image')
        cv2.setMouseCallback('image',self.mark)

        while True:
            cv2.imshow('image',self.img)
            k = cv2.waitKey(1) & 0xFF
            if k==ord('c'):
                self.cls_num +=1
                print("第{}类".format(self.cls_num+1))
            if k==27:
                break
        cv2.destroyAllWindows()
        result = cv2.cvtColor(self.markers,cv2.COLOR_BGR2GRAY)
        #不能用插值,会改变标签值
        result = result[::self.resize_rate,::self.resize_rate]
        cv2.imwrite('./markers.png',self.img)
        return result

3 主函数

首先获得标记
开运算
分水岭

要想分割的好，要避免不同区域的连通，相同区域的分裂。锐化和高斯模糊感觉用了效果不好。

img = cv2.imread('test.png',0)
#标记图像
DM = draw_markers(img,1)
markers = DM.draw()
#开运算
kernel = np.ones((5,5),dtype=np.uint8)
img = cv2.morphologyEx(img,cv2.MORPH_OPEN,kernel,iterations=1)
# img = cv2.erode(img,kernel=np.ones((5,5)),iterations=2)
# img = cv2.GaussianBlur(img,(5,5),sigmaX=0)

# #锐化
# kernel = np.array([[1, 1, 1], [1, -8, 1], [1, 1, 1]], dtype=np.float32)
# imgLaplacian = cv2.filter2D(img, cv2.CV_32F, kernel)
# sharp =np.float32(img)
# imgResult = sharp - imgLaplacian
# imgResult = np.clip(imgResult,0,255)
# imgResult = np.uint8(imgResult)
#开始分水岭算法
WSHED= watershed.WATERSHED(img,markers)
result = WSHED.run()
#将所有的分水岭转为255
result[result==-1]==255
result = result.astype(np.uint8)
imshow(result,'result.png')