  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.
