MRI_recon/code/Frequency-Diffusion/experiments/FSMNet/train_brats.py

import os
import sys
from tqdm import tqdm
from tensorboardX import SummaryWriter
import shutil
import argparse
import logging
import time
import torch
import numpy as np
import torch.optim as optim
from torchvision import transforms
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision.utils import make_grid
from networks.compare_models import build_model_from_name
from dataloaders.BRATS_dataloader_new import Hybrid as MyDataset
from dataloaders.BRATS_dataloader_new import RandomPadCrop, ToTensor, AddNoise
from networks.mynet import TwoBranch
from option import args
from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity


train_data_path = args.root_path
test_data_path = args.root_path
snapshot_path = "model/" + args.exp + "/"

os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
batch_size = args.batch_size * len(args.gpu.split(','))
max_iterations = args.max_iterations
base_lr = args.base_lr



def cc(img1, img2):
    eps = torch.finfo(torch.float32).eps
    """Correlation coefficient for (N, C, H, W) image; torch.float32 [0.,1.]."""
    N, C, _, _ = img1.shape
    img1 = img1.reshape(N, C, -1)
    img2 = img2.reshape(N, C, -1)
    img1 = img1 - img1.mean(dim=-1, keepdim=True)
    img2 = img2 - img2.mean(dim=-1, keepdim=True)
    cc = torch.sum(img1 * img2, dim=-1) / (eps + torch.sqrt(torch.sum(
        img1 **2, dim=-1)) * torch.sqrt(torch.sum(img2**2, dim=-1)))
    cc = torch.clamp(cc, -1., 1.)
    return cc.mean()



def gradient_calllback(network):
    for name, param in network.named_parameters():
        if param.grad is not None:
            # print("Gradient of {}: {}".format(name, param.grad.abs().mean()))

            if param.grad.abs().mean() == 0:
                print("Gradient of {} is 0".format(name))

        else:
            print("Gradient of {} is None".format(name))

class AMPLoss(nn.Module):
    def __init__(self):
        super(AMPLoss, self).__init__()
        self.cri = nn.L1Loss()

    def forward(self, x, y):
        x = torch.fft.rfft2(x, norm='backward')
        x_mag =  torch.abs(x)
        y = torch.fft.rfft2(y, norm='backward')
        y_mag = torch.abs(y)

        return self.cri(x_mag,y_mag)


class PhaLoss(nn.Module):
    def __init__(self):
        super(PhaLoss, self).__init__()
        self.cri = nn.L1Loss()

    def forward(self, x, y):
        x = torch.fft.rfft2(x, norm='backward')
        x_mag = torch.angle(x)
        y = torch.fft.rfft2(y, norm='backward')
        y_mag = torch.angle(y)

        return self.cri(x_mag, y_mag)

if __name__ == "__main__":
    ## make logger file
    if not os.path.exists(snapshot_path):
        os.makedirs(snapshot_path)

    logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO,
                        format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S')
    logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
    logging.info(str(args))

    network = TwoBranch(args).cuda()
    device = torch.device('cuda')
    network.to(device)

    if len(args.gpu.split(',')) > 1:
        network = nn.DataParallel(network)
        # network = nn.SyncBatchNorm.convert_sync_batchnorm(network)

    n_parameters = sum(p.numel() for p in network.parameters() if p.requires_grad)
    print('number of params: %.2f M' % (n_parameters / 1024 / 1024))

    db_train = MyDataset(split='train', MRIDOWN=args.MRIDOWN, SNR=args.low_field_SNR, 
                        #  transform=transforms.Compose([RandomPadCrop(), ToTensor(), AddNoise()]),
                         transform=transforms.Compose([RandomPadCrop(), ToTensor()]),
                         base_dir=train_data_path, input_normalize = args.input_normalize)
    
    db_test = MyDataset(split='test', MRIDOWN=args.MRIDOWN, SNR=args.low_field_SNR, 
                        transform=transforms.Compose([ToTensor()]),
                        base_dir=test_data_path, input_normalize = args.input_normalize)

    trainloader = DataLoader(db_train, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
    fixtrainloader = DataLoader(db_train, batch_size=1, shuffle=False, num_workers=4, pin_memory=True)
    testloader = DataLoader(db_test, batch_size=1, shuffle=False, num_workers=4, pin_memory=True)

    if args.phase == 'train':
        network.train()

        params = list(network.parameters())
        optimizer1 = optim.AdamW(params, lr=base_lr, betas=(0.9, 0.999), weight_decay=1e-4)
        scheduler1 = optim.lr_scheduler.StepLR(optimizer1, step_size=20000, gamma=0.5)

        writer = SummaryWriter(snapshot_path + '/log')

        iter_num = 0
        max_epoch = max_iterations // len(trainloader) + 1


        best_status = {'T1_NMSE': 10000000, 'T1_PSNR': 0, 'T1_SSIM': 0,
                       'T2_NMSE': 10000000, 'T2_PSNR': 0, 'T2_SSIM': 0}
        fft_weight=0.01
        criterion = nn.L1Loss().to(device, non_blocking=True)
        amploss = AMPLoss().to(device, non_blocking=True)
        phaloss = PhaLoss().to(device, non_blocking=True)
        start_time = time.time()

        for epoch_num in tqdm(range(max_epoch), ncols=70):
            time1 = time.time()
            debug_time = False

            # Data Preparation Time:  0.01880049705505371
            # Network Forward Time:  0.08233189582824707
            # Loss Calculation Time:  0.08654212951660156
            # Optimizer Step Time:  0.4485752582550049

            for i_batch, (sampled_batch, sample_stats) in enumerate(trainloader):
                time2 = time.time()
   
                t1_in, t1, t2_in, t2 = sampled_batch['image_in'].cuda(), sampled_batch['image'].cuda(), \
                                       sampled_batch['target_in'].cuda(), sampled_batch['target'].cuda()
                
                t1_krecon, t2_krecon = sampled_batch['image_krecon'].cuda(), sampled_batch['target_krecon'].cuda()
                

                time3 = time.time()

                if debug_time:
                    print("Data Preparation Time: ", time3 - time2)
                    print("t1, t2=", t1.shape, t2.shape)

                outputs = network(t2_in, t1_in)
                if debug_time:
                    print("Network Forward Time: ", time.time() - time2)

                loss = criterion(outputs['img_out'], t2) + \
                        fft_weight * amploss(outputs['img_fre'], t2) + fft_weight * phaloss(
                        outputs['img_fre'],
                        t2) + \
                        criterion(outputs['img_fre'], t2)
                if debug_time:
                    print("Loss Calculation Time: ", time.time() - time2)

                time4 = time.time()
                optimizer1.zero_grad()
                loss.backward()

                if args.clip_grad == "True":
                    ### clip the gradients to a small range.
                    torch.nn.utils.clip_grad_norm_(network.parameters(), 0.01)

                optimizer1.step()
                scheduler1.step()
                if debug_time:
                    print("Optimizer Step Time: ", time.time() - time2)

                time5 = time.time()

                # summary
                iter_num = iter_num + 1
                # writer.add_scalar('lr', scheduler1.get_lr(), iter_num)
                # writer.add_scalar('loss/loss', loss, iter_num)

                if iter_num % 100 == 0:
                    logging.info('iteration %d [%.2f sec]: learning rate : %f loss : %f ' % (iter_num, time.time()-start_time, scheduler1.get_lr()[0], loss.item()))
    
                if iter_num % 20000 == 0:
                    save_mode_path = os.path.join(snapshot_path, 'iter_' + str(iter_num) + '.pth')
                    torch.save({'network': network.state_dict()}, save_mode_path)
                    logging.info("save model to {}".format(save_mode_path))

                if iter_num > max_iterations:
                    break
                time1 = time.time()
            
            
            ## ================ Evaluate ================
            logging.info(f'Epoch {epoch_num} Evaluation:')
            # print()
            t1_MSE_all, t1_PSNR_all, t1_SSIM_all = [], [], []
            t2_MSE_all, t2_PSNR_all, t2_SSIM_all = [], [], []

            t1_MSE_krecon, t1_PSNR_krecon, t1_SSIM_krecon = [], [], []
            t2_MSE_krecon, t2_PSNR_krecon, t2_SSIM_krecon = [], [], []

            for (sampled_batch, sample_stats) in testloader:
                
                t1_in, t1, t2_in, t2 = sampled_batch['image_in'].cuda(), sampled_batch['image'].cuda(), \
                                       sampled_batch['target_in'].cuda(), sampled_batch['target'].cuda()
                
                t1_krecon, t2_krecon = sampled_batch['image_krecon'].cuda(), sampled_batch['target_krecon'].cuda()
                t_merge = torch.cat([t1_in, t2_in], dim=1)

                t2_out = network(t2_in, t1_in)['img_out']
                t1_out = None

                if args.input_normalize == "mean_std":
                    t1_mean = sample_stats['t1_mean'].data.cpu().numpy()[0]
                    t1_std = sample_stats['t1_std'].data.cpu().numpy()[0]
                    t2_mean = sample_stats['t2_mean'].data.cpu().numpy()[0]
                    t2_std = sample_stats['t2_std'].data.cpu().numpy()[0]

                    if t1_out is not None:
                        t1_img = (np.clip(t1.data.cpu().numpy()[0, 0] * t1_std + t1_mean, 0, 1) * 255).astype(np.uint8)
                        t1_out_img = (np.clip(t1_out.data.cpu().numpy()[0, 0] * t1_std + t1_mean, 0, 1) * 255).astype(np.uint8)
                        t1_krecon_img = (np.clip(t1_krecon.data.cpu().numpy()[0, 0] * t1_std + t1_mean, 0, 1) * 255).astype(np.uint8)

 
                    t2_img = (np.clip(t2.data.cpu().numpy()[0, 0] * t2_std + t2_mean, 0, 1) * 255).astype(np.uint8)
                    t2_out_img = (np.clip(t2_out.data.cpu().numpy()[0, 0] * t2_std + t2_mean, 0, 1) * 255).astype(np.uint8)
                    t2_krecon_img = (np.clip(t2_krecon.data.cpu().numpy()[0, 0] * t2_std + t2_mean, 0, 1) * 255).astype(np.uint8)

                else:
                    if t1_out is not None:
                        t1_img = (np.clip(t1.data.cpu().numpy()[0, 0], 0, 1) * 255).astype(np.uint8)
                        t1_out_img = (np.clip(t1_out.data.cpu().numpy()[0, 0], 0, 1) * 255).astype(np.uint8)
                        t1_krecon_img = (np.clip(t1_krecon.data.cpu().numpy()[0, 0], 0, 1) * 255).astype(np.uint8)
 
                    t2_img = (np.clip(t2.data.cpu().numpy()[0, 0], 0, 1) * 255).astype(np.uint8)
                    t2_out_img = (np.clip(t2_out.data.cpu().numpy()[0, 0], 0, 1) * 255).astype(np.uint8)
                    t2_krecon_img = (np.clip(t2_krecon.data.cpu().numpy()[0, 0], 0, 1) * 255).astype(np.uint8)


                if t1_out is not None:

                    MSE = mean_squared_error(t1_img, t1_out_img)
                    PSNR = peak_signal_noise_ratio(t1_img, t1_out_img)
                    SSIM = structural_similarity(t1_img, t1_out_img)
                    t1_MSE_all.append(MSE)
                    t1_PSNR_all.append(PSNR)
                    t1_SSIM_all.append(SSIM)

                    MSE = mean_squared_error(t1_img, t1_krecon_img)
                    PSNR = peak_signal_noise_ratio(t1_img, t1_krecon_img)
                    SSIM = structural_similarity(t1_img, t1_krecon_img)
                    t1_MSE_krecon.append(MSE)
                    t1_PSNR_krecon.append(PSNR)
                    t1_SSIM_krecon.append(SSIM)


                if t2_out is not None:
                    MSE = mean_squared_error(t2_img, t2_out_img)
                    PSNR = peak_signal_noise_ratio(t2_img, t2_out_img)
                    SSIM = structural_similarity(t2_img, t2_out_img)
                    t2_MSE_all.append(MSE)
                    t2_PSNR_all.append(PSNR)
                    t2_SSIM_all.append(SSIM)
                    # print("[t2 MRI] MSE:", MSE, "PSNR:", PSNR, "SSIM:", SSIM)
                    
                    MSE = mean_squared_error(t2_img, t2_krecon_img)
                    PSNR = peak_signal_noise_ratio(t2_img, t2_krecon_img)
                    SSIM = structural_similarity(t2_img, t2_krecon_img)
                    t2_MSE_krecon.append(MSE)
                    t2_PSNR_krecon.append(PSNR)
                    t2_SSIM_krecon.append(SSIM)

            if t1_out is not None:
                t1_mse = np.array(t1_MSE_all).mean()
                t1_psnr = np.array(t1_PSNR_all).mean()
                t1_ssim = np.array(t1_SSIM_all).mean()

                t1_krecon_mse = np.array(t1_MSE_krecon).mean()
                t1_krecon_psnr = np.array(t1_PSNR_krecon).mean()
                t1_krecon_ssim = np.array(t1_SSIM_krecon).mean()

            t2_mse = np.array(t2_MSE_all).mean()
            t2_psnr = np.array(t2_PSNR_all).mean()
            t2_ssim = np.array(t2_SSIM_all).mean()
                
            t2_krecon_mse = np.array(t2_MSE_krecon).mean()
            t2_krecon_psnr = np.array(t2_PSNR_krecon).mean()
            t2_krecon_ssim = np.array(t2_SSIM_krecon).mean()


            if t2_psnr > best_status['T2_PSNR']:
                best_status = {'T2_NMSE': t2_mse, 'T2_PSNR': t2_psnr, 'T2_SSIM': t2_ssim}

                best_checkpoint_path = os.path.join(snapshot_path, 'best_checkpoint.pth')

                torch.save({'network': network.state_dict()}, best_checkpoint_path)
                print('New Best Network:')

            logging.info(f"[T2 MRI:] average MSE: {t2_mse} average PSNR: {t2_psnr} average SSIM: {t2_ssim}")

            if iter_num > max_iterations:
                break
        print(best_status)
        save_mode_path = os.path.join(snapshot_path, 'iter_' + str(max_iterations) + '.pth')
        torch.save({'network': network.state_dict()},
                   save_mode_path)
        logging.info("save model to {}".format(save_mode_path))
        writer.close()