ABN

Anonymous repository for Anti-Blur Neural Networks for Multi-Stage Deformable Image Registration (ABN). This repository contains the implementation of ABN.

File listing

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.

Instructions on training ABN

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

Result

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