OctreeNCA: Single-Pass 184 MP Segmentation on Consumer Hardware

TL;DR

OctreeNCA introduces a VRAM-efficient, single-pass segmentation method for large medical images and videos, leveraging an octree-based neighborhood to enable global context understanding on consumer hardware.

Contribution

We propose OctreeNCA, a novel octree-based neighborhood generalization for NCAs, enabling efficient, global-aware segmentation of large inputs with significantly reduced VRAM usage.

Findings

Segments 184 MP pathology slices in one pass

Reduces VRAM usage by 90% compared to UNet

Enables real-time segmentation of high-resolution videos

Abstract

Medical applications demand segmentation of large inputs, like prostate MRIs, pathology slices, or videos of surgery. These inputs should ideally be inferred at once to provide the model with proper spatial or temporal context. When segmenting large inputs, the VRAM consumption of the GPU becomes the bottleneck. Architectures like UNets or Vision Transformers scale very poorly in VRAM consumption, resulting in patch- or frame-wise approaches that compromise global consistency and inference speed. The lightweight Neural Cellular Automaton (NCA) is a bio-inspired model that is by construction size-invariant. However, due to its local-only communication rules, it lacks global knowledge. We propose OctreeNCA by generalizing the neighborhood definition using an octree data structure. Our generalized neighborhood definition enables the efficient traversal of global knowledge. Since deep learning frameworks are mainly developed for large multi-layer networks, their implementation does not fully leverage the advantages of NCAs. We implement an NCA inference function in CUDA that further reduces VRAM demands and increases inference speed. Our OctreeNCA segments high-resolution images and videos quickly while occupying 90% less VRAM than a UNet during evaluation. This allows us to segment 184 Megapixel pathology slices or 1-minute surgical videos at once.