UnWave-Net: Unrolled Wavelet Network for Compton Tomography Image Reconstruction
This repository contains the official PyTorch implementation of the paper UnWave-Net: Unrolled Wavelet Network for Compton Tomography Image Reconstruction, presented at MICCAI 2024.
UnWave-Net is designed for image reconstruction in Compton tomography, utilizing unrolled optimization techniques in combination with wavelet transforms. The model excels at reconstructing images with reduced noise by incorporating a wavelet-based regularization term.
The UnWave-Net model has been trained and evaluated using the AAPM Compton Scattering Tomography Dataset. This dataset is essential for benchmarking the performance of Compton tomography reconstruction techniques.
- Source: Kaggle - AAPM Compton Scattering Tomography Dataset
- Description: The dataset consists of raw images and data obtained from Compton scattering tomography experiments. It includes data collected using different numbers of detectors, allowing for the evaluation of reconstruction performance under varying imaging conditions.
- Download: You can download the dataset directly from here.
We provide pre-trained models for UnWave-Net, trained on different datasets with and without noise. You can download these checkpoint files and use them to evaluate or fine-tune the models.
The table below lists the available checkpoints:
| Dataset | T | Noise | Checkpoint Size | Download Link |
|---|---|---|---|---|
| D_K = 100 | 16 | Free | 98.2 MB | Download |
| D_K = 150 | 16 | Free | 98.2 MB | Download |
| D_K = 100 | 16 | With | 98.2 MB | Download |
| D_K = 150 | 16 | With | 98.2 MB | Download |
To use these models, simply download the checkpoint files and load them using the provided evaluation scripts.
To use the dataset with UnWave-Net, follow these steps:
1. Clone this repo
git clone https://github.com/Ishak96/UnWave-Net.git
cd UnWave-Net2. Mamba installation If you don't have Conda installed, you can either install it manually or opt for Mamba, which is faster.
To install Mamba, follow this link to Mambaforge. Below are the latest installers with Mamba in the base environment:
| Operating System | Architecture | Download Link |
|---|---|---|
| Linux | x86_64 (amd64) | Mambaforge-Linux-x86_64 |
| macOS | x86_64 | Mambaforge-MacOSX-x86_64 |
| macOS | arm64 (Apple Silicon) | Mambaforge-MacOSX-arm64 |
| Windows | x86_64 | Mambaforge-Windows-x86_64 |
After installing Mambaforge, verify that it is functioning correctly by running the following command in your terminal:
mamba --version3. Install the environment
conda env create -f env.yml4. Activate the environment
conda activate UnWaveNet5. Download and extract the AAPM Compton Scattering Tomography Dataset from Kaggle.
6. Download the forward and adjoint operators and store them in a folder named radon. For each D_K, create a separate folder in radon; for example, D_K = 100:
mkdir radon
mkdir radon/100The table below lists the available forward and adjoint operators:
| D_K | Forward Operator | Adjoint Operator |
|---|---|---|
| 50 | Download | Download |
| 100 | Download | Download |
| 150 | Download | Download |
7. Generate the JSON file containing paths to individual images / collected data:
cd util
python generate_json_cst.py --data_dir PATH-TO-CST-MAYO-DATANow you are ready to train or test UnWave-Net! 😊
1. Training: To train the model, run the following command:
./scripts/run.sh2. Testing: To test the model with a specific checkpoint, use the command:
./scripts/run.sh default_hp.resume_from_checkpoint="path/to/checkpoint.ckpt" train=False test=TrueFeel free to check the run.sh file and modify the arguments as needed.
UnWave-Net is licensed under the MIT License. For details, please refer to the LICENSE file.
If you find our paper useful, please cite
@inproceedings{ishak2024unwave,
title={{UnWave-Net}: Unrolled Wavelet Network for {Compton Tomography} Image Reconstruction},
author={Ishak Ayad and Cécilia Tarpau and Javier Cebeiro and Maï K. Nguyen},
booktitle={Medical Image Computing and Computer-Assisted Intervention--MICCAI 2024},
year={2024},
organization={Springer}
}- This work was granted access to the HPC resources of IDRIS under the allocation 2021-[AD011012741] / 2022-[AD011013915] provided by GENCI and supported by DIM Math Innov funding.
- The dataset used in this work, the AAPM Low-Dose CT Grand Challenge Dataset, was provided by the American Association of Physicists in Medicine (AAPM) and is available at AAPM Grand Challenge.