基于阈值与watershed的三维点簇状信号分割

三维成像数据中的物体分割，不能简单的先逐层分析每个slice，再堆叠回去。因此这里再专门提供一个三维点簇状信号分割的函数。

之前已经介绍过怎么用 watershed 分水岭算法来解决点簇状信号的分割问题，对于三维数据来说，思路是一样的，而且得益于我调用的函数能够支持高维数据的处理，所以稍作修改就能实现对三维成像数据的实例分割。

效果如下：

具体代码如下：

1
import matplotlib.pyplot as plt
2
import matplotlib as mpl
3
mpl.rcParams['font.family'] = 'DengXian'
4
plt.rcParams['axes.unicode_minus'] = False
5

6

7
from skimage import io
8
import numpy as np
9
from scipy import ndimage as ndi
10
from skimage.segmentation import watershed
11
from skimage.feature import peak_local_max
12
from skimage import filters
13
from skimage.color import label2rgb
14

15

16

17

18
def wat_seg_3d(volume,
19
               threshold,
20
               peak_min_distance=3):
21
    '''
22
    volume: 三维图像 (NumPy array, typically Z, Y, X order, but depends on your data)
23
    peak_min_distance: 决定了搜索peak的范围，在每个维度上，例如 3 表示在 3x3x3 的邻域内
24
    返回一张label图像，背景为0，目标已按region id 进行赋值, 且为uint16
25
    '''
26

27

28
    # 图像二值化，使用在二维分割或目视检查的自定义threshold
29
    binary = volume > threshold
30

31

32
    # 距离变换 (distance_transform_edt works on N-dim)
33
    # It computes the Euclidean distance from non-zero (foreground) pixels to zero (background) pixels.
34
    distance = ndi.distance_transform_edt(binary)
35

36

37
    # 准备标记点 (peak_local_max works on N-dim)
38
    # Finds local maxima in the original intensity volume as potential markers
39
    coords = peak_local_max(volume,
40
                            min_distance=peak_min_distance,
41
                            threshold_abs=threshold, # Use the determined threshold
42
                            exclude_border=False) # Decide if peaks on borders should be considered
43

44

45
    # Initialize markers array with zeros (undetermined)
46
    markers = np.zeros(shape=volume.shape, dtype=np.int32)
47

48

49
    # 将局部极大值位置设置为前景标记点
50
    # 我们从标签 2 开始标记前景（标签 1 保留给背景，标签 0 是未确定区域）
51
    # 使用 enumerate 为每个局部极大值赋予唯一的标签
52
    # coords will be an array of shape (N, 3) for 3D, where N is the number of peaks
53
    for i, coord in enumerate(coords):
54
        # coord is a tuple/list of indices (z, y, x) for 3D
55
        # Need to use tuple indexing for N-dim arrays
56
        markers[tuple(coord)] = i + 2 # Labels start from 2, 3, 4...
57

58

59
    # 确定“确定无疑的背景”区域，直接设定阈值 (np.where works on N-dim)
60
    inds = np.where(volume <= threshold)
61
    markers[inds] = 1 # 将确定背景区域设置为标签 1
62

63

64
    # 应用分水岭算法 (watershed works on N-dim)
65
    # Use the negative distance map as the 'surface' for the watershed
66
    labels = watershed(-distance, markers)
67

68

69
    # Output formatting: shift labels so background is 0
70
    # The watershed background label is typically 1 if provided as a marker.
71
    # Subtracting 1 makes marker 1 -> label 0 (background), marker 2 -> label 1, etc.
72
    rlabel = labels - 1
73

74

75
    # Ensure output type is uint16
76
    rlabel = rlabel.astype('uint16')
77

78

79
    return rlabel