Implementation of and experiments based on the QSMNet (https://arxiv.org/abs/1803.05627).
Now CAD-QSMNet is not an implementation of the original QSMNet any more, but we indeed heavily drew inspiration from it.
This work was done to submit to the QSM Reconstruction Challenge 2.0.
Names are listed in alphabetical order:
- Jaewon Chung (@jaywonchung)
- Jinho Park (@jinh0park)
- Yunchan Hwang(@yunchfruit)
Special thanks to Woojin Jung (@wjjung93) for his guidance and helpful advice during the challenge.
-
Clone this repository.
git clone https://github.com/jinh0park/QSMFirstOrDeath.git
-
Install prerequisites.
Using Anaconda3.
torch >= 1.1.0- Cuda version == 8.0
conda install pytorch=0.4.1 cuda80 -c pytorch pip install -r requirements.txt
- Cuda version == 9.0
conda install pytorch torchvision cudatoolkit=9.0 -c pytorch pip install -r requirements.txt
- Cuda version == 10.0
conda install pytorch torchvision cudatoolkit=10.0 -c pytorch pip3 install -r requirements.txt
-
Pull the latest changes from the repository.
git pull
-
Create a new branch in the repository.
git checkout -b your_branch_name
-
Change parameters in
constants.py. -
Change the code directly.
data_utils.pydefines how the code loads data from .mat files.model.pydefines the QSMNet model layers and their connections.train.pydefines the training logic. You can change the optimizer here.loss.pydefines the loss functions.train.pycalls thetotal_lossfunction for the final loss.
-
Start training from scratch with
python3 train.py
Or, resume training from a saved training state with
python3 train.py "path/to/saved/file"FYI,
train.pysaves the training state (epochs done, model weights, optimizer, and learning rate scheduler) every epoch (by default). -
Check your training status with PyTorch TensorBoard.
tensorboard --logdir runs --host localhost
Once tensorboard runs, navigate to the given address (e.g. localhost:6006) on your machine's web browser.
-
When you're done, you may merge the branch to
masterif it should be part of the baseline model by sending a pull request. If not, you can delete the branch withgit checkout master git branch -d your_branch_name
But be careful; this will remove all the changes you made to that branch!
inference.py helps you perform inferences with the checkpoint file. It creates .mat and .nii files with inference results on the challenge data and the test set.
- Requirements
- The checkpoint file and its path in the form
Checkpoints/<model_architecture>/<experience_name>/<model_name>.ckpt. - The
config.jsonfile for the checkpoint, in theCheckpoints/<model_architecture>/<experience_name>folder. - The training data used to train the model. We need the input and output mean/std from it.
- Run
bash ../QSMFirstOrDeath/prepare_inference.sh /path/to/NAS/list3in theDatafolder. This will (hopefully) have you meet all the requirements below this line. - In the
../Datafolder:Challenge_stage1_sim1_fix.matChallenge_stage1_sim2_fix.matChallenge_stage1_ROI_mask.matTest_0718.matTest_ROI_mask.mat
- In the
../Data/Conventionalfolder:Sim1_Input.niiandSim2_Input.niiSim1_{method}.niiandSim2_{method}.nii, wheremethodis the reconstruction method.Test_{method}.matwheremethodis the reconstruction method.- Also check and edit the two lists named '
methods' ininference.pyappropriately. The first one is for the challenge inference, and the second one is for the test set inference.
- The checkpoint file and its path in the form
- Usage
python inference.py 'Checkpoints/<model_architecture>/<experience_name>/<model_name>.ckpt' - A quick inference function:
inf- Import:
from inference import inf
- Arguments:
checkpoint_file: path to the checkpoint file, relative to the QSMFirstOrDeath folderfile: (optional) network input of extension.mat,.niior.nii.gznumpy: (optional) network input of typenp.ndarraytensor: (optional) network input of typetorch.Tensordevice: (optional) which device to use for inference. If not given, takes value fromconstants.py.
- Usage: Provide one, two, or all of
file,numpy, ortensor.- When only one input argument is given,
infreturns its inference value. Forfileandnumpy, the output is anp.ndarray. Fortensor, the output is atorch.Tensor. - When more than one of the input arguments are given,
infreturns a tuple of outputs. Inferences are made separately.
- When only one input argument is given,
- Import: