Unsupervised Feature Loss (UFLoss) for High Fidelity Deep learning (DL)-based reconstruction
Official github repository for the paper High Fidelity Deep Learning-based MRI Reconstruction with Instance-wise Discriminative Feature Matching Loss. In this work, a novel patch-based Unsupervised Feature loss (UFLoss) is proposed and incorporated into the training of DL-based reconstruction frameworks in order to preserve perceptual similarity and high-order statistics. In-vivo experiments indicate that adding the UFLoss encourages sharper edges with higher overall image quality under DL-based reconstruction framework. Our implementations are in PyTorch
Installation
To use this package, install the required python packages (tested with python 3.8 on Ubuntu 20.04 LTS):
pip install -r requirements.txt
Dataset
We used a subset of FastMRI knee dataset for the training and evaluation. We used E-SPIRiT to pre-compute sensitivity maps using BART. Post-processed data (including Sens Maps, Coil combined images) and pre-trained model can be requested by emailing [email protected].
Update We provide our data-preprocessing code at UFloss_training/data_preprocessing.py. This script computes the sensitivity maps and performs data normalization and coil combination. BART toolbox is required for computing the sensitivity maps. Follow the installation instructions on the website and add the following lines to your .bashrc file.
export PYTHONPATH="${PYTHONPATH}:/python/ " export PATH=": $PATH "
To run the data-preprocessing code, download and unzip the fastMRI Multi-coil knee dataset. Simplu run
python data_preprocessing.py -l <path to your fastMRI multi-coil dataset> -t <target directory> -c <size for your E-SPIRiT calibration region>
Step 0: Patch Extraction
To extract patches from the fully-smapled training data, go to the UFloss_training/ folder and run patch_extraction.py to extract patches. Please specify the directories of the training dataset and the target folder. Instructions are avaible by runing:
python patch_extraction.py -h
Step 1: Train the UFLoss feature mapping network
To train the UFLoss feature mapping network, go to the UFloss_training/ folder and run patch_learning.py. We provide a demo training script to perform the training on fully-sampled patches:
bash launch_training_patch_learning.sh
Visualiztion (Patch retrival results, shown below) script will be available soon.
Step 2: Train the DL-based reconstruction with UFLoss
To train the DL-based reconstruction with UFLoss, we provide our source code here at DL_Recon_UFLoss/. We adoped MoDL as our DL-based reconstruction network. We provide training scripts for MoDL with and without UFLoss at DL_Recon_UFLoss/models/unrolled2D/scripts:
bash launch_training_MoDL_traditional_UFLoss_256_demo.sh
You can easily paly around with the parameters by editing the training script. One representative reconstruction results is shown as below.
Perform inference with the trained model
To perform the inference reconstruction on the testing set, we provide an inference script at DL_Recon_UFLoss/models/unrolled2D/inference_ufloss.py. run the following command for inference:
python inference_ufloss.py --data-path <Path to the dataset>
--device-num <Which device to train on>
--exp-dir <Path where the results should be saved>
--checkpoint <Path to an existing checkpoint>
Acknoledgements
Reconstruction code borrows heavily from fastMRI Github repo and DL-ESPIRiT by Christopher Sandino. This work is a colaboration between UC Berkeley and GE Healthcare. Please contact [email protected] if you have any questions.
Citation
If you find this code useful for your research, please consider citing our paper High Fidelity Deep Learning-based MRI Reconstruction with Instance-wise Discriminative Feature Matching Loss:
@article{wang2021high,
title={High Fidelity Deep Learning-based MRI Reconstruction with Instance-wise Discriminative Feature Matching Loss},
author={Wang, Ke and Tamir, Jonathan I and De Goyeneche, Alfredo and Wollner, Uri and Brada, Rafi and Yu, Stella and Lustig, Michael},
journal={arXiv preprint arXiv:2108.12460},
year={2021}
}
75 Dec 19, 2022
16 Nov 4, 2020
287 Jan 6, 2023
1 Jan 18, 2022
51 Dec 13, 2022
114 Dec 10, 2022
31 Dec 8, 2022
170 Jan 4, 2023
57 Jan 5, 2023
371 Dec 30, 2022
1 Feb 7, 2022
157 Jan 1, 2023
77 Dec 21, 2022
117 Dec 28, 2022
32 Sep 21, 2022
71 Dec 4, 2022