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