The efficacy of segmentation algorithms is frequently compromised by topological errors like overlapping regions, disrupted connections, and voids. To tackle this problem, we introduce a novel loss function, namely Topology-Aware Focal Loss (TAFL), that incorporates the conventional Focal Loss with a topological constraint term based on the Wasserstein distance between the ground truth and predicted segmentation masks' persistence diagrams. By enforcing identical topology as the ground truth, the topological constraint can effectively resolve topological errors, while Focal Loss tackles class imbalance. We begin by constructing persistence diagrams from filtered cubical complexes of the ground truth and predicted segmentation masks. We subsequently utilize the Sinkhorn-Knopp algorithm to determine the optimal transport plan between the two persistence diagrams. The resultant transport plan minimizes the cost of transporting mass from one distribution to the other and provides a mapping between the points in the two persistence diagrams. We then compute the Wasserstein distance based on this travel plan to measure the topological dissimilarity between the ground truth and predicted masks. We evaluate our approach by training a 3D U-Net with the MICCAI Brain Tumor Segmentation (BraTS) challenge validation dataset, which requires accurate segmentation of 3D MRI scans that integrate various modalities for the precise identification and tracking of malignant brain tumors. Then, we demonstrate that the quality of segmentation performance is enhanced by regularizing the focal loss through the addition of a topological constraint as a penalty term.
翻译:分割算法的有效性常因拓扑错误(如重叠区域、连接中断和空洞)而受损。为解决此问题,我们提出一种新型损失函数——拓扑感知焦点损失(TAFL),该函数将传统焦点损失与基于真实标签和预测分割掩膜持续图之间Wasserstein距离的拓扑约束项相结合。通过强制拓扑结构与真实标签一致,拓扑约束可有效解决拓扑错误,而焦点损失则处理类别不平衡问题。我们首先从真实标签和预测分割掩膜的滤复形复合结构构建持续图,随后利用Sinkhorn-Knopp算法确定两个持续图之间的最优传输方案。该传输方案使质量从一种分布转移到另一种分布的成本最小化,并提供两个持续图中点之间的映射关系。基于此传输方案,我们计算Wasserstein距离以衡量真实标签与预测掩膜之间的拓扑差异。我们通过使用MICCAI脑肿瘤分割(BraTS)挑战验证数据集训练三维U-Net来评估本方法,该数据集要求对整合多种模态的三维MRI扫描进行精确分割,以精准识别和追踪恶性脑肿瘤。实验结果表明,通过将拓扑约束作为惩罚项加入焦点损失进行正则化,可提升分割性能质量。