Promptable cancer segmentation using minimal expert-curated data

TL;DR

The paper tackles the high cost of pixel-level cancer annotations for segmentation by introducing a data-efficient promptable segmentation method for prostate cancer in MR images. It combines a weakly supervised and a fully supervised classifier to guide a point-based prompt through a spiral crop search to a full ROI segmentation, trained on only 8 weak labels and 24 full segmentations. The approach outperforms existing promptable methods and approaches fully supervised performance while using far less data, demonstrating practical potential for clinical deployment and generalization to other scarce-label domains. An open-source implementation and extensive ablations underscore the method's efficiency and transferability to other cancer localization tasks.

Abstract

Automated segmentation of cancer on medical images can aid targeted diagnostic and therapeutic procedures. However, its adoption is limited by the high cost of expert annotations required for training and inter-observer variability in datasets. While weakly-supervised methods mitigate some challenges, using binary histology labels for training as opposed to requiring full segmentation, they require large paired datasets of histology and images, which are difficult to curate. Similarly, promptable segmentation aims to allow segmentation with no re-training for new tasks at inference, however, existing models perform poorly on pathological regions, again necessitating large datasets for training. In this work we propose a novel approach for promptable segmentation requiring only 24 fully-segmented images, supplemented by 8 weakly-labelled images, for training. Curating this minimal data to a high standard is relatively feasible and thus issues with the cost and variability of obtaining labels can be mitigated. By leveraging two classifiers, one weakly-supervised and one fully-supervised, our method refines segmentation through a guided search process initiated by a single-point prompt. Our approach outperforms existing promptable segmentation methods, and performs comparably with fully-supervised methods, for the task of prostate cancer segmentation, while using substantially less annotated data (up to 100X less). This enables promptable segmentation with very minimal labelled data, such that the labels can be curated to a very high standard.

Figures (4)

• Figure 1: The training (top) and inference (bottom) pipelines for the weakly-supervised and fully-supervised classifiers.
• Figure 2: The point prompt leads to a spiral search where each crop is scored by the two classifiers to determine whether to mark a crop as positive or negative. Combining all positive crops gives the final segmentation.
• Figure 3: Plot of crop size $(w, h, d)$ against dice score. $(w, h)$ on the bottom axis and $d$ on the top axis.
• Figure 4: Samples of MR scans. Ground truth is in blue, predicted is in red.