PancreasSeg-UNet (CT/CBCT)

This repository contains the trained weights for a two-stage 3D pancreas segmentation model designed for both CT and CBCT abdominal scans. The model combines a lightweight UNet for global localization and an Attention UNet for fine-grained segmentation within the predicted bounding box, achieving robust performance even in challenging cases with low contrast and ambiguous boundaries.

Example patient – GT (left) vs. prediction (right), Dice = 0.66

Example patient – GT (left) vs. prediction (right), Dice = 0.73

Example patient – GT (left) vs. prediction (right), Dice = 0.84

Model Overview

Attribute	Description
Task	Automatic pancreas segmentation from 3D CT / CBCT volumes
Input	Single‑channel 3D image (normalized HU range → [0,1])
Output	Binary pancreas mask (same spatial dimensions as input)
Architecture	Two‑stage: UNet (localization) + Attention UNet (refinement)
Framework	PyTorch + MONAI
Training Data	Pancreatic CT/CBCT dataset from The Cancer Imaging Archive
Preprocessing	Resampling to (1.5×1.5×3.0 mm), intensity clipping to [-150,250] HU, z‑score normalisation, pancreas‑centered cropping
Loss Function	DiceCELoss (Dice + Cross‑Entropy)
Optimiser	Adam (lr=1e-3) + Warmup Cosine scheduler
Hardware	NVIDIA GPU (CUDA)

Architecture Details

The model consists of two sub‑networks that are trained jointly but used in a sequential inference pipeline:

Stage 1 – Localisation UNet

• Type: Standard 3D UNet (monai.networks.nets.UNet)
• Channels: (16, 32, 64, 128) with strides (2,2,2)
• Input size: 128×128×64 (resized full volume)
• Purpose: Coarse pancreas localisation → produces a bounding box around the predicted region.

Stage 2 – Refinement Attention UNet

• Type: 3D Attention UNet (monai.networks.nets.AttentionUnet)
• Channels: (32, 64, 128, 256, 512) with strides (2,2,2,2)
• Input size: 128×128×96 (crop from Stage 1 bbox)
• Purpose: Fine segmentation using attention gates to focus on relevant anatomical structures.
• Gradient checkpointing is enabled during training to reduce memory consumption.

The two stages are trained simultaneously:

• Loss = loss_stage2 + 0.5 * loss_stage1
• Stage 2 crops are generated online with RandCropByPosNegLabeld (positive/negative ratio 1:2) and strong augmentation (rotation, flip, intensity shift).

Dataset

The model was trained on the Pancreatic CT/CBCT Segmentation collection available from The Cancer Imaging Archive (TCIA).
Please cite the original dataset when using this model:

Hong, J., Reyngold, M., Crane, C., Cuaron, J., Hajj, C., Mann, J., Zinovoy, M., Yorke, E., LoCastro, E., Apte, A. P., & Mageras, G. (2021). Breath-hold CT and cone-beam CT images with expert manual organ-at-risk segmentations from radiation treatments of locally advanced pancreatic cancer [Data set]. The Cancer Imaging Archive. https://doi.org/10.7937/TCIA.ESHQ-4D90

Inference Pipeline

During inference, the following steps are performed:

1. Normalise the input volume (clamp to [-150,250] HU, scale to [0,1]).
2. Resize the whole volume to 128×128×64 and pass it through Stage 1 UNet.
3. Extract bounding box from the coarse prediction (add margin of 30 voxels, then clamp to original dimensions).
4. Crop the original‑resolution volume to the bounding box.
5. Apply sliding‑window inference (window size 128×128×96, stride 4) with Stage 2 Attention UNet.
6. Place the fine prediction back into a full‑sized zero array to obtain the final mask.

Data Preprocessing

1. DICOM / RTSTRUCT parsing – All DICOM series are indexed; RTSTRUCT files are matched to their corresponding CT series using SeriesInstanceUID.
2. Resampling – CT images are resampled to a common spacing (1.5 mm, 1.5 mm, 3.0 mm) using linear interpolation; masks use nearest neighbour.
3. Intensity standardisation – Clip to [-150, 250] HU (pancreas window) and linearly rescale to [0, 1].
4. Bounding box extraction – From the ground truth mask, a margin of 30 voxels is added; each volume is cropped to that box and saved as ct.npy and mask.npy.
5. Dataset split – Patient‑wise random split: 28 training, 6 validation, rest test (fixed random seed 42).

Code & Usage

The full training and inference code, along with a Streamlit‑based graphical user interface for running the model on your own CT/CBCT volumes, is available in the GitHub repository.

License

Mozilla Public License Version 2.0 - Feel free to use and modify