rNCA: Self-Repairing Segmentation Masks

TL;DR

This paper introduces rNCA, a neural cellular automaton-based method that refines imperfect segmentation masks by iteratively repairing topological errors, leading to significant improvements across various medical imaging tasks.

Contribution

The paper demonstrates how Neural Cellular Automata can be repurposed as a universal, learnable refinement mechanism for correcting topological errors in segmentation masks.

Findings

rNCA improves Dice and clDice scores by 2-3%.

Reduces Betti errors by 60% ($eta_0$) and 20% ($eta_1$).

Repairs 61.5% of broken myocardium masks in zero-shot setting.

Abstract

Accurately predicting topologically correct masks remains a difficult task for general segmentation models, which often produce fragmented or disconnected outputs. Fixing these artifacts typically requires hand-crafted refinement rules or architectures specialized to a particular task. Here, we show that Neural Cellular Automata (NCA) can be directly re-purposed as an effective refinement mechanism, using local, iterative updates guided by image context to repair segmentation masks. By training on imperfect masks and ground truths, the automaton learns the structural properties of the target shape while relying solely on local information. When applied to coarse, globally predicted masks, the learned dynamics progressively reconnect broken regions, prune loose fragments and converge towards stable, topologically consistent results. We show how refinement NCA (rNCA) can be easily applied to repair common topological errors produced by different base segmentation models and tasks: for fragmented retinal vessels, it yields 2-3% gains in Dice/clDice and improves Betti errors, reducing $\beta_0$ errors by 60% and $\beta_1$ by 20%; for myocardium, it repairs 61.5% of broken cases in a zero-shot setting while lowering ASSD and HD by 19% and 16%, respectively. This showcases NCA as effective and broadly applicable refiners.