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

John Kalkhof, Camila González, Anirban Mukhopadhyay

https://arxiv.org/pdf/2302.03473.pdf

Med-NCA performs segmentation with only 70k parameters by iterating the same local rule over each cell of an image. @Hessian_AI, @CS_TUDarmstadt, @ipmi2023, @anirbanakash pic.twitter.com/N3umQy7Q0h

— John Kalkhof (@kalkjo) March 21, 2023

<script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script>

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.

To get started with this repository simply follow these few steps:

Quickstart

1. Install requirements of repository: pip install -r requirements.txt
2. Download prostate dataset from: http://medicaldecathlon.com/
3. Adapt img_path and label_path in train_Med_NCA.ipynb
4. Run train_Med_NCA.ipynb
5. To view results in tensorboard: tensorboard --logdir path