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.py
defines how the code loads data from .mat files.model.py
defines the QSMNet model layers and their connections.train.py
defines the training logic. You can change the optimizer here.loss.py
defines the loss functions.train.py
calls thetotal_loss
function 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.py
saves 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
master
if 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.json
file 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/list3
in theData
folder. This will (hopefully) have you meet all the requirements below this line. - In the
../Data
folder:Challenge_stage1_sim1_fix.mat
Challenge_stage1_sim2_fix.mat
Challenge_stage1_ROI_mask.mat
Test_0718.mat
Test_ROI_mask.mat
- In the
../Data/Conventional
folder:Sim1_Input.nii
andSim2_Input.nii
Sim1_{method}.nii
andSim2_{method}.nii
, wheremethod
is the reconstruction method.Test_{method}.mat
wheremethod
is the reconstruction method.- Also check and edit the two lists named '
methods
' ininference.py
appropriately. 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
,.nii
or.nii.gz
numpy
: (optional) network input of typenp.ndarray
tensor
: (optional) network input of typetorch.Tensor
device
: (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,
inf
returns its inference value. Forfile
andnumpy
, the output is anp.ndarray
. Fortensor
, the output is atorch.Tensor
. - When more than one of the input arguments are given,
inf
returns a tuple of outputs. Inferences are made separately.
- When only one input argument is given,
- Import: