Med-NCA: Robust and Lightweight Segmentation with Neural Cellular Automata

TL;DR

Med-NCA introduces a lightweight, robust neural cellular automata approach for high-resolution medical image segmentation, outperforming UNet while being highly resource-efficient and suitable for deployment on low-power devices.

Contribution

The paper presents Med-NCA, a novel end-to-end training pipeline for high-resolution medical image segmentation using neural cellular automata, addressing VRAM constraints and convergence issues.

Findings

Outperforms UNet by 2-3% Dice on hippocampus and prostate segmentation.

Is 500 times smaller than traditional models.

Operates effectively on low-resource devices like Raspberry Pi.

Abstract

Access to the proper infrastructure is critical when performing medical image segmentation with Deep Learning. This requirement makes it difficult to run state-of-the-art segmentation models in resource-constrained scenarios like primary care facilities in rural areas and during crises. The recently emerging field of Neural Cellular Automata (NCA) has shown that locally interacting one-cell models can achieve competitive results in tasks such as image generation or segmentations in low-resolution inputs. However, they are constrained by high VRAM requirements and the difficulty of reaching convergence for high-resolution images. To counteract these limitations we propose Med-NCA, an end-to-end NCA training pipeline for high-resolution image segmentation. Our method follows a two-step process. Global knowledge is first communicated between cells across the downscaled image. Following that, patch-based segmentation is performed. Our proposed Med-NCA outperforms the classic UNet by 2% and 3% Dice for hippocampus and prostate segmentation, respectively, while also being 500 times smaller. We also show that Med-NCA is by design invariant with respect to image scale, shape and translation, experiencing only slight performance degradation even with strong shifts; and is robust against MRI acquisition artefacts. Med-NCA enables high-resolution medical image segmentation even on a Raspberry Pi B+, arguably the smallest device able to run PyTorch and that can be powered by a standard power bank.