如何高性能实现3D类别分割掩码的孔洞填充操作？

我们拥有多类别3D分割掩码，每个类别对应独立的标签ID，需要对每个类别的分割区域做孔洞填充处理，填充规则为仅填充被单一类别完全包围的孔洞区域，向边界开放的区域不填充，不同类别之间的背景0值保持不变。

示例输入输出 #

输入矩阵 #

[
  [
    [ 1, 1, 1, 2, 2, 2 ],
    [ 1, 1, 1, 2, 2, 2 ],
    [ 1, 1, 1, 2, 2, 2 ],
    [ 0, 3, 0, 0, 4, 0 ],
    [ 3, 3, 3, 4, 0, 4 ],
    [ 0, 3, 0, 0, 4, 0 ],
  ],
  [
    [ 1, 1, 1, 2, 2, 2 ],
    [ 1, 0, 1, 2, 0, 0 ],
    [ 1, 1, 1, 2, 2, 2 ],
    [ 0, 3, 0, 0, 4, 0 ],
    [ 3, 0, 3, 4, 0, 4 ],
    [ 0, 3, 0, 0, 4, 0 ],
  ],
  [
    [ 1, 1, 1, 2, 2, 2 ],
    [ 1, 1, 1, 2, 2, 2 ],
    [ 1, 1, 1, 2, 2, 2 ],
    [ 0, 3, 0, 0, 4, 0 ],
    [ 3, 3, 3, 4, 4, 4 ],
    [ 0, 3, 0, 0, 4, 0 ],
  ],
]

期望输出 #

[
  [
    [ 1, 1, 1, 2, 2, 2 ],
    [ 1, 1, 1, 2, 2, 2 ],
    [ 1, 1, 1, 2, 2, 2 ],
    [ 0, 3, 0, 0, 4, 0 ],
    [ 3, 3, 3, 4, 0, 4 ],
    [ 0, 3, 0, 0, 4, 0 ],
  ],
  [
    [ 1, 1, 1, 2, 2, 2 ],
    [ 1, 1, 1, 2, 0, 0 ],
    [ 1, 1, 1, 2, 2, 2 ],
    [ 0, 3, 0, 0, 4, 0 ],
    [ 3, 3, 3, 4, 0, 4 ],
    [ 0, 3, 0, 0, 4, 0 ],
  ],
  [
    [ 1, 1, 1, 2, 2, 2 ],
    [ 1, 1, 1, 2, 2, 2 ],
    [ 1, 1, 1, 2, 2, 2 ],
    [ 0, 3, 0, 0, 4, 0 ],
    [ 3, 3, 3, 4, 4, 4 ],
    [ 0, 3, 0, 0, 4, 0 ],
  ],
]

示例说明：仅中间切片的1类、3类内部孔洞被填充，2类区域向侧边开放、4类区域向背部开放因此不做处理，类别间的0值保持不变。

现有实现代码 #

目前基于scipy.ndimage.morphology接口及numpy实现了三个性能最优的版本，代码如下：

import numpy as np
from scipy.ndimage.morphology import binary_fill_holes, label, generate_binary_structure, binary_dilation

def fill_holes6(img: np.ndarray, applied_labels: np.ndarray) -> np.ndarray:
    output = np.zeros_like(img)
    for i in applied_labels:
        output[binary_fill_holes(img == i)] = i

    return output

def fill_holes7(img: np.ndarray, applied_labels: np.ndarray) -> np.ndarray:
    output = np.zeros(img.shape, dtype=int)
    for i in applied_labels:
        tmp = np.zeros(img.shape, dtype=bool)
        binary_dilation(tmp, structure=None, iterations=-1, mask=img != i, origin=0, border_value=1, output=tmp)
        output[np.logical_not(tmp)] = i
        
    return output

# 新增方法：
def fill_holes8(img: np.ndarray, applied_labels: np.ndarray) -> np.ndarray:
    connectivity = 1
    footprint = generate_binary_structure(img.ndim, connectivity)
    background_mask = img == 0
    components, num_components = label(background_mask, structure=footprint)
    filled_holes = np.zeros_like(img)
    for component_label in range(1, num_components + 1):
        component_mask = components == component_label
        component_neighborhood = np.pad(img, 1, constant_values=-1)[binary_dilation(np.pad(component_mask, 1), structure=footprint)]

        neighbor_labels = np.unique(component_neighborhood)
        if len(neighbor_labels) == 2 and -1 not in neighbor_labels:
            neighbor_label = neighbor_labels[1]
            filled_holes[component_mask] = neighbor_label

    return img + filled_holes

性能测试与结果 #

测试代码（匹配真实业务数据分布）：

import time
import pandas as pd

def measure(funs, t):
    res = []
    for _ in range(t):
        ra = np.random.randint(10, 40)
        sh = np.random.randint(200, 400, 3)
        img = np.random.randint(0, ra, sh)

        applied_labels = np.unique(img)[1:]

        fun_res = []
        for fun in funs:
            start = time.time()
            fun(img, applied_labels)
            end = time.time()
            fun_res.append(end - start)
        res.append(fun_res)
    return np.min(res, axis=0), np.max(res, axis=0), np.mean(res, axis=0), np.std(res, axis=0)

print(measure([fill_holes6, fill_holes7], t=10))

首批实现版本在t=100的测试下执行耗时如下：

现存问题：现有实现性能极低，最新的fill_holes7版本仅比fill_holes3快1.27倍，无法满足业务需求。请问是否存在更高效的孔洞填充实现方案？

补充说明 #

我已先向scipy项目提交功能需求，被引导先到stackoverflow咨询，也向MONAI项目提交了功能需求，同时提交了基于迭代腐蚀的fill_holes7实现的PR，对应文档见monai.transforms.FillHoles；此外还实现了基于连通域标记(CCL)的fill_holes8版本。MONAI社区欢迎任何优化该方法性能的PR，可前往对应仓库提交issue和PR。

内容的提问来源于stack exchange，提问作者Spenhouet