Super-resolution of biomedical volumes with 2D supervision

CVPR Workshop on Computer Vision for Microscopy Image Analysis (CVMI) 2024

Website / arXiv / MLiNS Lab / OpenSRH / Model checkpoint / Sample volumes

Installation

1. Clone MSDSR github repo

   git clone [email protected]:MLNeurosurg/msdsr.git

2. Install miniconda: follow instructions here

3. Create conda environment:

   conda create -n msdsr python=3.11

4. Activate conda environment:

   conda activate msdsr

5. Install package and dependencies

   pip install -e .

Model checkpoint and sample volumes

We release our pretrained DDPM model checkpoint and sample 3D volumes. They are available at the links below:

Model checkpoint / Sample volumes

Training / evaluation instructions

The code base is written using PyTorch Lightning, with custom network and datasets for OpenSRH.

To train MSDSR on the OpenSRH dataset:

1. Download OpenSRH - request data here.
2. Update the sample config file in train/config/train_msdsr.yaml with desired configurations.
3. Change directory to train and activate the conda virtual environment.
4. Use train/train_msdsr.py to start training:

   python train_msdsr.py -c=config/train_msdsr.yaml

To evaluate with your trained model:

1. Update the sample config files in eval/config/*.yaml with the checkpoint path and other desired configurations per file. If you are using the released checkpoint, place the checkpoint in the path $log_dir/$exp_name/msdsr_cvmi24/models/d17986ac.ckpt.
2. Change directory to eval and activate the conda virtual environment.
3. Use the evaluation scripts in eval/*.py for evaluation. For example:

   # paired 2D evaluation
   python eval_paired.py -c=config/eval_paired.yaml
   
   # generate 3D volumes evaluation
   python generate_volumes.py -c=config/generate_volumes.yaml
   
   # generate metrics for 3D volumes [require paired 2D evaluation and 3D volume generated]
   python compute_volume_metrics.py -c=config/compute_metrics.yaml