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
