Numpy实现分水岭分割算法【未完结】

from queue import PriorityQueue
import numpy as np


class Pixel(object):
    """
    像素信息，包含像素的坐标和梯度
    """

    def __init__(self, gradient, location):
        self.gradient = gradient
        self.location = location

    def __lt__(self, other):
        return self.gradient < other.gradient

    def __str__(self):
        return f'{self.gradient}:{self.location}'


class Watershed:
    grad: np.ndarray
    markers: np.ndarray

    def __init__(self, grad, markers=None):
        """
        分水岭分割算法：\n
        （1）基于标记控制的分水岭算法：原始梯度图 + 标记矩阵\n
        （2）基于标记控制和强制最小技术的分水岭算法：强制最小后的梯度图

        :param grad:    输入的梯度图像（用于优先级排序：梯度小的优先级大）
        :param markers: 输入的标记矩阵（用于区分汇水区：汇水区分界处为分水岭）
        """
        # 进入队列的标记
        self.VISITED = -2
        # 分水岭的标记
        self.W_SHED = -1
        # 未知点的标记
        self.UNKNOWN = 0
        # 像素队列（优先级为梯度值，梯度小的优先级高：类似于汇水区汇水，由低至高涨水）
        self.queue = PriorityQueue()
        # 输入梯度图像的拷贝
        self.grad = grad.copy()
        # 输入标记矩阵
        if markers is None:
            self.markers = np.zeros_like(grad)
        else:
            self.markers = markers.copy()
        # opencv的视觉边缘设置为分水岭（东南西北）
        self.markers[:, self.markers.shape[1] - 1] = self.W_SHED
        self.markers[0, :] = self.W_SHED
        self.markers[:, 0] = self.W_SHED
        self.markers[self.markers.shape[0] - 1, :] = self.W_SHED
        # 定义4邻域规则，邻域内有无已知前景区域或者背景区域
        self.neighbourhood4 = lambda x, y: {(x - 1, y): self.markers[x - 1, y],
                                            (x + 1, y): self.markers[x + 1, y],
                                            (x, y - 1): self.markers[x, y - 1],
                                            (x, y + 1): self.markers[x, y + 1]}
        self.nbh4_any_known = lambda x, y: np.any(np.array([val for loc, val in self.neighbourhood4(x, y).items()]) > 0)

    def ws_push(self, location: tuple) -> None:
        """
        将传入的像素点的位置信息入队，并在markers上做标记

        :param location:    输入的像素点位置信息（x, y）
        :return:            None
        """
        gradient: int = self.grad[location[0], location[1]]
        pixel = Pixel(gradient, location)
        self.queue.put(pixel)
        self.markers[location[0], location[1]] = self.VISITED

    def ws_pop(self) -> tuple:
        """
        将队首像素点信息出队

        :return:    队首像素点信息（若队为空则返回[-1,-1]）
        """
        return self.queue.get().location if self.queue.qsize() > 0 else (-1, -1)

    def label_nbh4_pixels(self, location: tuple) -> None:
        """
        在markers中，为每一个像素点对应的位置处设置label

        :param location:    输入的像素点位置
        :return:            None
        """
        (x, y), label = location, self.UNKNOWN
        for nbh4_loc, nbh4_val in self.neighbourhood4(x, y).items():
            # 当邻域内像素是前景区域或者背景区域时
            if nbh4_val > 0:
                # 如果是第一个邻域点，则让label赋值为该邻域点的值
                if label == self.UNKNOWN:
                    label = nbh4_val
                # 如果该邻域点的值不等于之前领域点的值，则说明pixel是分水岭（前景和背景的交界处，也就是边缘处）
                elif label != nbh4_val:
                    label = self.W_SHED
                # 如果该邻域点的值与之前邻域点的值是相等的，则保持不变
                else:
                    pass
        # 将该点在markers中标记为label
        self.markers[x, y] = label

    def push_nbh4_pixels(self, location: tuple) -> None:
        """
        对各个像素点进行四领域分析：\n
            如果满足以下条件则入队：\n
            （1）pixel既不是分水岭，也没有入队（没有归属于积水池）；\n
            （2）邻域点既不是前景区域，也不是背景区域，也就是说它属于unknown区域；

        :param location:    输入的像素坐标
        :return:            None
        """
        (x, y) = location
        # 首先得确保该像素点不是分水岭，如果是分水岭就没必要再入队寻找了（因为已经找到了）
        if self.markers[x, y] != self.W_SHED:
            for nbh4_loc, nbh4_val in self.neighbourhood4(x, y).items():
                # 只对未知区域进行扩散（已知区域已经确定了，就没必要扩散了）
                if nbh4_val == self.UNKNOWN:
                    self.ws_push(nbh4_loc)

    def watershed(self) -> np.ndarray:
        """
        执行基于标记的分水岭分割算法

        :return:    处理后的标记矩阵
        """
        # 将markers的初始点（属于unknown区域且4邻域内存在已知区域）放入优先队列
        for row in range(1, self.markers.shape[0] - 1):
            for col in range(1, self.markers.shape[1] - 1):
                # 保证先从unknown区域的边缘处开始寻找分水岭
                if self.markers[row, col] == self.UNKNOWN and self.nbh4_any_known(row, col):
                    self.ws_push((row, col))
        # 将优先队列中的像素点位置信息“先出队后入队”，并保证每个像素点有且仅有一次处理机会
        counter, max_its = 0, self.markers.shape[0] * self.markers.shape[1]
        while self.queue.qsize() > 0 and counter < max_its:
            # 将梯度最高的像素点出队（梯度越高的点也可能是分水岭）
            location = self.ws_pop()
            # 判断该像素点是否为分水岭（依据为4邻域内是否同时存在不同类型的已知区域）
            self.label_nbh4_pixels(location)
            # 如果该点不是分水岭，则将其4邻域内unknown区域的像素点入队
            self.push_nbh4_pixels(location)
            counter += 1
        # 返回处理后的标记矩阵
        return self.markers