Point-Unet: A Context-aware Point-based Neural Network for Volumetric Segmentation

TL;DR

The paper tackles the memory bottlenecks of voxel-based 3D volumetric segmentation by introducing Point-Unet, a three-module point-based framework. It combines saliency attention for ROI estimation, context-aware sampling for single-pass ROI-focused point cloud generation, and a point-based segmentation network with a Dilated Residual Block encoder to produce voxel-level segmentation results. Training uses Generalized Dice Loss to address class imbalance, and experiments on Pancreas and BraTS show competitive or superior performance with reduced memory and faster inference compared to voxel-based baselines. The results establish a scalable, efficient alternative for 3D medical image segmentation with potential for broader clinical deployment.

Abstract

Medical image analysis using deep learning has recently been prevalent, showing great performance for various downstream tasks including medical image segmentation and its sibling, volumetric image segmentation. Particularly, a typical volumetric segmentation network strongly relies on a voxel grid representation which treats volumetric data as a stack of individual voxel `slices', which allows learning to segment a voxel grid to be as straightforward as extending existing image-based segmentation networks to the 3D domain. However, using a voxel grid representation requires a large memory footprint, expensive test-time and limiting the scalability of the solutions. In this paper, we propose Point-Unet, a novel method that incorporates the efficiency of deep learning with 3D point clouds into volumetric segmentation. Our key idea is to first predict the regions of interest in the volume by learning an attentional probability map, which is then used for sampling the volume into a sparse point cloud that is subsequently segmented using a point-based neural network. We have conducted the experiments on the medical volumetric segmentation task with both a small-scale dataset Pancreas and large-scale datasets BraTS18, BraTS19, and BraTS20 challenges. A comprehensive benchmark on different metrics has shown that our context-aware Point-Unet robustly outperforms the SOTA voxel-based networks at both accuracies, memory usage during training, and time consumption during testing. Our code is available at https://github.com/VinAIResearch/Point-Unet.

Figures (5)

• Figure 1: (a): a 2D voxel grid and a segmentation rendered by volume rendering; (b): a PC from the grid and the point-based segmentation results.
• Figure 2: Point-Unet takes a volume as input and consists of 3 modules: saliency attention network, context-aware sampling and point segmenation network.
• Figure 3: Our proposed saliency attention network.
• Figure 4: Our proposed Point-Unet segmentation on volumetric data.
• Figure 5: Performance analysis. (a) With a single inference, our Point-Unet outperforms RandLA-Net, which requires multiple iterations at inference. (b) Memory requirement for training with batch size 1 and inference time with difference volume patch sizes.