rNCA: Self-Repairing Segmentation Masks

2025-12-15

AI Summary

General-purpose segmentation models often produce masks with topological errors (e.g., disconnections, fragmentation). Existing correction methods rely on handcrafted rules or task-specific architectures. This paper introduces the first plug-and-play Neural Cellular Automaton (NCA) module for topology-aware mask refinement: taking coarse segmentation outputs as input, it performs local neighborhood updates guided by image context and evolves masks end-to-end via iterative, adaptive refinement—yielding connected and topologically consistent results. The module requires no task-specific design or manual priors and enables zero-shot cross-model and cross-task transfer. On retinal vessel segmentation, it improves Dice and clDice by 2–3% and reduces β₀ and β₁ topological errors by 60% and 20%, respectively. In myocardial segmentation, it repairs 61.5% of disconnected cases in zero-shot fashion, reducing ASSD and Hausdorff Distance by 19% and 16%.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Deep Generative Models & Autoencoders
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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 $β_0$ errors by 60% and $β_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.

Problem

• Repair fragmented or disconnected segmentation masks
• Improve topological correctness of segmentation outputs
• Fix common topological errors across different segmentation tasks

Innovation

• Neural Cellular Automata refines segmentation masks locally
• Learns structural properties from imperfect masks and ground truths
• Repairs topological errors across different segmentation tasks