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.

翻译：分割算法的有效性常因拓扑错误（如区域重叠、连接中断及空洞）而受损。为解决此问题，我们提出一种新型损失函数——拓扑感知聚焦损失（Topology-Aware Focal Loss, TAFL），该函数将传统聚焦损失与基于真实分割掩码与预测分割掩码持久化图之间Wasserstein距离的拓扑约束项相结合。通过强制预测结果与真实掩码保持相同拓扑结构，拓扑约束能够有效解决拓扑错误，而聚焦损失则处理类别不平衡问题。我们首先从真实掩码与预测掩码的滤波立方复形中构建持久化图，随后利用Sinkhorn-Knopp算法确定两个持久化图之间的最优传输方案。该传输方案在最小化质量从一个分布传输到另一分布的成本的同时，提供两持久化图各点之间的映射关系。基于此传输方案计算Wasserstein距离，以衡量真实掩码与预测掩码之间的拓扑差异性。我们采用MICCAI脑肿瘤分割挑战（BraTS）验证数据集训练三维U-Net，该数据集要求对整合多种模态的三维MRI扫描进行精确分割，以精准识别和追踪恶性脑肿瘤。实验证明，通过将拓扑约束作为惩罚项加入聚焦损失进行正则化，分割性能质量得到显著提升。