Heteroscedastic Uncertainty Estimation for Probabilistic Unsupervised Registration of Noisy Medical Images
├── Dataset
| ├── ACDC # Place the downloaded dataset here
| | ├── train
| | ├── val
| | ├── test
| ├── CAMUS
| | ├── train
| | ├── ...
├── Code
| ├── hetero_uncertainty
| | ├── src
| | | ├── voxelmorph
| | | ├── transmorph
| | ├── ...
├── Models
| ├── hetero_uncertainty
| | ├── vxm_ACDC_motion.pt
| | ├── vxm_ACDC_variance.pt
| | ├── ...
├── Logs
conda create -f requirements.yml -n hetero_env # Create an environment named hetero_env
conda activate hetero_env
Please download ACDC dataset [website] and CAMUS dataset [website] from their websites. The pretrained weights can be found at [here]. We also provide a sample script, preprocess.ipynb, to assist with preprocessing the raw dataset.
train.py: The script used to train the model. Can be found at src/voxelmorph/train.py and src/transmorph/train.py
-
--dataset: Specifies the dataset to be used for training. Choose from['ACDC', 'CAMUS']. -
--bidir: Enables bidirectional training, which simultaneously optimizes forward and backward transformations between image pairs. -
--motion-loss-type: Defines the type of motion loss used during training.'wmse'indicates a weighted mean squared error loss. -
--variance-loss-type: Specifies the type of variance-related loss.'beta-NLL'refers to a beta-weighted negative log-likelihood loss. -
--batch-size: Batch size. -
--accumulation_steps: Enables gradient accumulation, effectively simulating a larger batch size without increasing memory usage. -
--epochs: Specifies the number of training epochs. -
--warm_start: Initializes the network by sequentially training specific components. -
--warm_start_epoch: Number of epochs to sequentially train each component. -
--model-dir: Path to save the model. Example:'../../Models/hetero_uncertainty/ACDC/...' -
--lambda: Sets the regularization loss weight$\lambda$ . -
--use_wandb: Activates logging with Weights & Biases (WandB) for tracking experiments. -
--wandb-project: Specifies the WandB project name. -
--wandb-name: Sets a specific run name for the experiment in WandB, making it easier to identify in the dashboard.
test.py: The script used to evaluate the trained model. Example scripts can be found at here.
-
--dataset: Specifies the dataset to be used for testing. Choose from['ACDC', 'CAMUS']. -
--test-dir: Path to the directory containing test data. -
--result-dir: Path to save the results of the testing process. -
--model-motion: Path to the saved motion estimator model. -
--model-variance: Path to the saved variance estimator model. -
--inshape: Specifies the input shape for the model, given as dimensions (e.g.,128 128).
We evaluated our approach using a contour-based metric by warping the ED (end-diastolic) myocardium segmentation with the estimated displacement to reconstruct the ES (end-systolic) segmentation. The reconstructed ES segmentation was then compared with the ground truth ES segmentation provided by the dataset. It is important to note that no segmentation labels were used during training. The quantitative results were reproduced by running inference with the pre-trained weights, as demonstrated in evaluation.ipynb.
To quantitatively assess our estimated heteroscedastic uncertainty for non-uniform noise, we utilized sparsification error metrics and computed the area under the sparsification error curve. Detailed explanations and results can be found in evaluation.ipynb.
@inproceedings{zhang2024heteroscedastic,
title={Heteroscedastic Uncertainty Estimation Framework for Unsupervised Registration},
author={Zhang, Xiaoran and Pak, Daniel H and Ahn, Shawn S and Li, Xiaoxiao and You, Chenyu and Staib, Lawrence H and Sinusas, Albert J and Wong, Alex and Duncan, James S},
booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
pages={651--661},
year={2024},
organization={Springer}
}
We use implementation of Voxelmorph, Transmorph.