Anonymous repository for Anti-Blur Neural Networks for Multi-Stage Deformable Image Registration (ABN). This repository contains the implementation of ABN.
The corresponding folder contains the source code for 2D and 3D image registration.
-
2D : 2D image registration task
- dataset : Sample dataset for 2D image registration
- main.py : Main code for ABN training
- model.py : Supporting models
- train.py : Supporting training
- utils.py : Supporting functions
-
3D : 3D image registration task
- dataset : Sample dataset for 3D image registration
- main.py : Main code for ABN training
- model.py : Supporting models
- preprocess_data.py : Using for preprocess data
- train.py : Supporting training
- utils.py : Supporting functions
Note that all public datasets used in the paper can be found here:
- FFHQ : for 2D face image.
- LPBA40 : for 3D brain MRI.
- Mindboggle101 : for 3D brain MRI.
The following script is for training:
python main.py
Parameters:
-
2D :
- img_size : size of input image, default 64
- num_stage : number of stages of ABN, default 10
- train_set_name : file name of training dataset, default 2D_face_train.npy
- test_set_name : file name of test dataset, default 2D_face_test.npy
- batch_size : batch size, default 16
- num_epochs : number of epochs, default 5000
- model_name : model, default ABN_2D
- loss_name : loss function, default MSE
- smooth_name : regularization, default second_Grad
- learning_rate : learning rate, default 0.0001
- lamda : lamda of regularization, default 10
- penalty : penalty of regularization, default l2
- save_every_epoch : gap of saving model, default 20
- num_sample : number of images displayed, default l6
- num_lines : number of grid lines displayed in the warped image, default 15
- grid_sample_orig : warping image from the original source image, default True
-
3D :
- img_size : size of input image, default 96
- num_stage : number of stages of ABN, default 10
- train_set_name : file name of training dataset, default LPBA40_train_sub.npy
- test_set_name : file name of test dataset, default LPBA40_test_sub.npy
- batch_size : batch size, default 1
- num_epochs : number of epochs, default 1000
- model_name : model, default ABN_3D
- loss_name : loss function, default NCC
- smooth_name : regularization, default second_Grad
- learning_rate : learning rate, default 0.0001
- lamda : lamda of regularization, default 10
- penalty : penalty of regularization, default l2
- save_every_epoch : gap of saving model, default 10
- sample_orig, : warping image from the original source image, default True
The results can be find after training.
- loss_log :
- model_name.txt : log file of the model
- model :
- model_name.pth : saved model
- sample_img :
- o : target images
- t : source images
- p_stage : warped images by stage
- p_stage_grid : warped images by stage with grid displayed