opt_grid_diff_ms.py

import numpy as np
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import torch.nn.functional as nnf
import argparse
import os
from torch import optim
import sys
import torch

from utils.csv_logger import infer_csv
from data import datasets
from models.dnvf.networks.grid import Grid
from models.dnvf.networks.siren import VecInt
import utils.losses as losses
import utils.utils as utils
from utils.train_utils import (adjust_learning_rate, 
                               save_checkpoint)



'''
parse command line arg
'''
parser = argparse.ArgumentParser()

parser.add_argument('--dataset', type=str, default='OASIS')
parser.add_argument('--save_dir', type=str, default='./infer_results/', help='The directory to save infer results')
parser.add_argument('--weight_jdet', type=int, default=50)
parser.add_argument('--val_iter', type=int, default=5)
parser.add_argument('--lr', type=float, default=1e-3)
parser.add_argument(
    "--downsample_factors",
    nargs="+",          # accept one or more arguments
    type=int,           # specify type
    required=True,
    help="list of the downsample factors. coarest to finest"
)
parser.add_argument(
    "--optimization_epochs",
    nargs="+",          # accept one or more arguments
    type=int,           # specify type
    required=True,
    help="list of the optimization epoches. coarest to finest"
)

args = parser.parse_args()

def main():
    csv_name = 'diff_grid_jDet{}_ms_{}_epoch_{}_{}'.format(args.weight_jdet, "_".join([str(d) for d in args.downsample_factors]), "_".join([str(e) for e in args.optimization_epochs]), args.dataset) + '.csv'
    
    downsample_factors = args.downsample_factors 
    optimization_epochs = args.optimization_epochs   
    val_iter = args.val_iter  
    weight_jdet = args.weight_jdet
    lr = args.lr

    save_dir = args.save_dir + args.dataset + '/' + 'grid' + '/'
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
        
    '''
    initialize dataset
    '''
    test_dataset, img_size = datasets.load_test_dataset(args.dataset)
    test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=0, pin_memory=False, drop_last=False)

    '''
    initialize spatial transformation function
    '''
    reg_model = utils.register_model(img_size, 'nearest')
    reg_model.cuda()
    reg_model_img = utils.register_model(img_size, 'bilinear')
    reg_model_img.cuda()
    vec_int = VecInt(img_size).cuda()
    
    '''
    initialize optimizer and loss functions
    '''
    criterion_ncc = losses.NCC_vxm()
    criterion_reg = losses.Grad3d(penalty='l2')
    criterion_jdet = losses.NegJDet(img_size).cuda()
    
    """start inferring"""    
    eval_dsc_def = utils.AverageMeter()
    eval_dsc_raw = utils.AverageMeter()
    eval_det = utils.AverageMeter()

    
    for data_idx, data in enumerate(test_loader):
    
        data = [t.cuda() for t in data]
        x, y, x_seg, y_seg = data
        loss_all = utils.AverageMeter()
        print('Processing pair: {}'.format(data_idx))
        
        for idx, (downsample_factor, optimization_epoch) in enumerate(zip(downsample_factors, optimization_epochs)):
            grid_input_shape = np.array(img_size) // downsample_factor
            if idx == 0:
                """Prepare Input Grid"""
                grid_input = torch.zeros(1, 3, grid_input_shape[0], grid_input_shape[1], grid_input_shape[2])
                print("flow field input shape: ", grid_input.shape)
            else:
                grid_input = nnf.interpolate(vec_field.flow.detach().clone(), size=list(grid_input_shape), mode='trilinear', align_corners=True)
            
            """Initialize model"""
            vec_field = Grid(grid_input)
            vec_field = vec_field.cuda()
            adam = optim.Adam(vec_field.parameters(), lr=lr, weight_decay=0, amsgrad=True)
            print('Current scale: {}x downsample, {} optimization epochs'.format(downsample_factor, optimization_epoch))
            
            for epoch in range(optimization_epoch):
                vec_field.train()
                adjust_learning_rate(adam, epoch, optimization_epoch, lr)
                output_norm_siren_vec = nnf.interpolate(vec_field.flow, size=img_size, mode='trilinear', align_corners=True)
                output_full_vec = output_norm_siren_vec
                output_full_vec[:, 0] = output_full_vec[:, 0] / 2 * (img_size[0] - 1)
                output_full_vec[:, 1] = output_full_vec[:, 1] / 2 * (img_size[1] - 1)
                output_full_vec[:, 2] = output_full_vec[:, 2] / 2 * (img_size[2] - 1)
                output_full_flow = vec_int(output_full_vec)
                
                warped_x = reg_model_img([x, output_full_flow])
                '''use ncc loss'''
                loss_ncc = criterion_ncc(warped_x, y)
                loss_reg = criterion_reg(output_full_flow)
                loss_jdet = criterion_jdet(output_full_flow)
                loss = loss_ncc + loss_reg + loss_jdet * weight_jdet
                loss_vals = [loss_ncc, loss_reg, loss_jdet]
                
                loss_all.update(loss.item(), y.numel())
                adam.zero_grad()
                loss.backward()
                adam.step()
                
                current_lr = adam.state_dict()['param_groups'][0]['lr']
                sys.stdout.write('\r Pair {}: Epoch [{}/{}] - loss {:.4f}, Img Sim: {:.6f}, Reg: {:.6f}, JDet: {:.6f}, lr: {:.6f}'.format(
                    data_idx, epoch, optimization_epoch, loss.item(), loss_vals[0].item(), loss_vals[1].item(), loss_vals[2].item(), current_lr))
                sys.stdout.flush()
                
                if epoch % val_iter == 0 or epoch == optimization_epoch - 1:
                    with torch.no_grad():
                        def_out = reg_model([x_seg.cuda().float(), output_full_flow.cuda()])
                        '''update DSC'''
                        if args.dataset == "OASIS":
                            dsc_trans = utils.dice_OASIS(def_out.long(), y_seg.long())
                            dsc_raw = utils.dice_OASIS(x_seg.long(), y_seg.long())
                        elif args.dataset == "IXI":
                            dsc_trans = utils.dice_IXI(def_out.long(), y_seg.long())
                            dsc_raw = utils.dice_IXI(x_seg.long(), y_seg.long())
                        elif args.dataset == "MindBoggle":
                            dsc_trans = utils.dice_MindBoggle(def_out.long(), y_seg.long())
                            dsc_raw = utils.dice_MindBoggle(x_seg.long(), y_seg.long())
                        elif args.dataset == "LPBA":
                            dsc_trans = utils.dice_LPBA(def_out.long(), y_seg.long())
                            dsc_raw = utils.dice_LPBA(x_seg.long(), y_seg.long())
                        elif args.dataset == "AbdomenCT":
                            dsc_trans = utils.dice_AbdomenCT(def_out.long(), y_seg.long())
                            dsc_raw = utils.dice_AbdomenCT(x_seg.long(), y_seg.long())
                        
                        '''update Jdet'''
                        jac_det = utils.jacobian_determinant_vxm(output_full_flow.detach().cpu().numpy()[0, :, :, :, :])
                        tar = y.detach().cpu().numpy()[0, 0, :, :, :]
                        Jdet = np.sum(jac_det <= 0) / np.prod(tar.shape)
                        
                        print('\n Val: Pair {}: Epoch [{}/{}] - DSC: {:.6f}, Jdet: {:.8f}, lr: {:.6f}\n'.format(
                    data_idx, epoch, optimization_epoch, dsc_trans, Jdet, current_lr))
                            
            torch.cuda.empty_cache()
            
        eval_det.update(Jdet, x.size(0))
        eval_dsc_raw.update(dsc_raw, x.size(0))
        eval_dsc_def.update(dsc_trans, x.size(0))
        infer_csv(save_dir, csv_name, data_idx, dsc_raw, dsc_trans, Jdet)
        print()
    infer_csv(save_dir, csv_name, 'avg', eval_dsc_raw.avg, eval_dsc_def.avg, eval_det.avg)
    infer_csv(save_dir, csv_name, 'std', eval_dsc_raw.std, eval_dsc_def.std, eval_det.std)




if __name__ == '__main__':
    main()