[Paper Review] Localise to segment: crop to improve organ at risk segmentation accuracy

Abraham George Smith, Denis Kutnár+3 more|arXiv (Cornell University)|Apr 10, 2023

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes a two-stage deep learning approach—localisation followed by organ segmentation—that improves organ-at-risk segmentation accuracy in radiotherapy planning. By cropping images to focus on relevant anatomy before segmentation, the method reduces class imbalance, accelerates training, and boosts Dice scores, especially for smaller organs like the spleen, pancreas, and heart.

ABSTRACT

Increased organ at risk segmentation accuracy is required to reduce cost and complications for patients receiving radiotherapy treatment. Some deep learning methods for the segmentation of organs at risk use a two stage process where a localisation network first crops an image to the relevant region and then a locally specialised network segments the cropped organ of interest. We investigate the accuracy improvements brought about by such a localisation stage by comparing to a single-stage baseline network trained on full resolution images. We find that localisation approaches can improve both training time and stability and a two stage process involving both a localisation and organ segmentation network provides a significant increase in segmentation accuracy for the spleen, pancreas and heart from the Medical Segmentation Decathlon dataset. We also observe increased benefits of localisation for smaller organs. Source code that recreates the main results is available at \href{https://github.com/Abe404/localise_to_segment}{this https URL}.

Motivation & Objective

• Improve segmentation accuracy for organs at risk in radiotherapy by reducing the impact of class imbalance in medical image datasets.
• Investigate whether a two-stage localisation-and-segmentation pipeline enhances performance compared to single-stage full-image segmentation.
• Evaluate the impact of localisation on training efficiency, convergence speed, and model stability across diverse organ sizes.
• Assess whether localisation provides greater benefits for smaller organs due to increased class imbalance.
• Determine if ground-truth-based localisation offers consistent advantages over end-to-end training on full-resolution images.

Proposed method

• Utilized a 3D U-Net architecture for both localisation and segmentation networks, with shared model structure to simplify implementation.
• Applied zero-padding in convolutional layers to maintain consistent input and output spatial dimensions during training.
• Used group normalisation instead of batch normalisation to stabilize training with small batch sizes (2 per GPU).
• Combined dice and cross-entropy loss to address class imbalance, particularly critical for small organs.
• Trained on cropped patches of size 64×256×256 to maximize context while fitting within 64GB GPU memory constraints.
• Implemented early stopping and model checkpointing to prevent overfitting and improve training efficiency.

Figure 1: Illustration showing the three different methods compared, including the baseline, two stage involving both a localisation network and organ segmentation network and the organ segmentation network that uses the ground truth to localise.

Experimental results

Research questions

• RQ1 Does a two-stage localisation-and-segmentation pipeline improve organ-at-risk segmentation accuracy compared to single-stage full-image training?
• RQ2 How does localisation affect training time, convergence rate, and model stability across different organ sizes?
• RQ3 Is the benefit of localisation more pronounced for smaller organs with higher class imbalance?
• RQ4 Can ground-truth bounding boxes for localisation consistently outperform end-to-end training on full images?
• RQ5 To what extent does localisation mitigate the negative effects of class imbalance in medical image segmentation?

Key findings

• The two-stage localisation approach significantly improved Dice scores for the spleen, pancreas, and heart, with the greatest gains observed in smaller organs.
• Training with localised crops reduced training time and improved convergence speed compared to full-image patch training, despite training two networks.
• The baseline single-stage model showed unstable training dynamics, including fluctuating Dice scores, likely due to severe class imbalance and inconsistent foreground presence in patches.
• The heart (left atrium) had the slowest convergence rate, attributed to both a small dataset and high class imbalance, though localisation improved performance.
• The pancreas, despite having the largest training set, showed only 85% convergence rate, indicating that dataset size alone does not overcome class imbalance challenges.
• Ground-truth-based localisation provided consistent improvements, suggesting manual or interactive localisation could be valuable in clinical workflows for high-accuracy segmentation.