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

from tqdm import tqdm
from tensorboardX import SummaryWriter
import logging, time, os, sys
import torch.optim as optim
from torchvision import transforms
from torch.utils.data import DataLoader
from dataloaders.BRATS_dataloader_new import Hybrid as MyDataset
from dataloaders.BRATS_dataloader_new import RandomPadCrop, ToTensor
from networks.mynet import TwoBranch
from utils.option import args
from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity

from utils.utils import *
from frequency_diffusion.degradation.k_degradation import get_ksu_kernel, apply_tofre, apply_to_spatial
from networks_time.mynet import DiffTwoBranch

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

# --use_time_model True --use_kspace True --ACCELERATIONS 4 --MRIDOWN 4X --low_field_SNR 20 --input_normalize mean_std
DEBUG = args.DEBUG
use_time_model = args.use_time_model
use_kspace = args.use_kspace     # PSNR: 30.138548551934974 average SSIM: 0.770964106980312
                      # PSNR: 31.325274490046855 average SSIM: 0.8589609042898623  4X if not

                      # PSNR: 29.846184815515585 average SSIM: 0.8797758188214125 -> 31.18, 0.77
                      # PSNR: 28.494279128317515 average SSIM: 0.8179950512965841  8X if not

# kspace_refine = True  # Albu with 41.33,  w/ 42.00
# mask_vacant = False
frequency_distortion = True



num_timesteps = args.num_timesteps  #30
image_size    = args.image_size     #240
distortion_sigma = 10/255

if args.MRIDOWN == "4X":
    accelerate_mask = np.load("./dataloaders/example_mask/brats_4X_mask.npy")
    accelerate_mask = torch.from_numpy(accelerate_mask).unsqueeze(0).clone().float()
else:
    accelerate_mask = None

# Output a list of k-space kernels
kspace_masks = get_ksu_kernel(num_timesteps, image_size,
                              ksu_routine="LogSamplingRate",
                              accelerated_factor=args.ACCELERATIONS[0],
                              accelerate_mask=accelerate_mask
                              )
np.save(f"./dataloaders/example_mask/brats_{args.ACCELERATIONS[0]}_kspace_mask.npy", kspace_masks)

kspace_masks = torch.from_numpy(np.asarray(kspace_masks)).cuda()



if __name__ == "__main__":
    ## make logger file
    if use_kspace:
        snapshot_path = snapshot_path.rstrip("/") + f'_t{num_timesteps}_kspace/'

    if use_time_model:
        snapshot_path = snapshot_path.rstrip("/") + '_time/'

    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))

    if use_time_model:
        network = DiffTwoBranch(args).cuda()
    else:
        network = TwoBranch(args).cuda()
    device  = torch.device('cuda')
    network.to(device)

    if len(args.gpu.split(',')) > 1:
        network = nn.DataParallel(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()]),
                         base_dir=train_data_path, input_normalize = args.input_normalize, use_kspace=use_kspace)
    
    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)
        if not use_kspace:
            scheduler1 = optim.lr_scheduler.StepLR(optimizer1, step_size=20000, gamma=0.5)
        else:
            scheduler1 = optim.lr_scheduler.StepLR(optimizer1, step_size=40000, gamma=0.5)

        writer = SummaryWriter(snapshot_path + '/log')

        iter_num  = 0
        max_epoch = max_iterations // len(trainloader) + 1
        if use_kspace:
            max_epoch = max_epoch * num_timesteps

        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)
        freloss   = Frequency_Loss().to(device, non_blocking=True)
        start_time = time.time()
        mask = None

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

            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()

                # Degradation
                if use_kspace:
                    b = t1_in.shape[0]
                    t = torch.randint(0, num_timesteps, (b,), device=device).long()
                    mask = kspace_masks[t]

                    target_fft, _ = apply_tofre(t2.clone(), mask)
                    fft, mask     = apply_tofre(t2_in.clone(), mask)

                    # if np.random.rand() > (1 / (1 + num_timesteps)):
                    fft = target_fft * mask + fft * (1 - mask)  # Seems too easy

                    # Frequency Noise
                    if frequency_distortion:
                        fft_magnitude = torch.abs(fft)     # 幅度
                        fft_phase     = torch.angle(fft)   # 相位

                        # Add noise to unmasked frequencies to maintain the stochasticity that diffusion models typically rely on.
                        sigma = distortion_sigma * torch.abs(torch.randn(1)).item()
                        noise = torch.randn_like(fft_magnitude) * sigma
                        noise_magnitude = noise * fft_magnitude * mask # + noise * (1 - mask)
                        fft_magnitude += noise_magnitude

                        sigma = distortion_sigma / 2 * torch.abs(torch.randn(1)).item()
                        noise = torch.randn_like(fft_phase) * sigma
                        noise_pha = noise * fft_phase * mask # + noise * (1 - mask)
                        fft_phase += noise_pha

                        fft = fft_magnitude * torch.exp(1j * fft_phase)

                    t2_in = apply_to_spatial(fft)

                time3 = time.time()

                if use_time_model and use_kspace:
                    outputs = network(t2_in, t1, t)
                else:
                    outputs = network(t2_in, t1)

                loss =  criterion(outputs['img_out'], t2) + criterion(outputs['img_fre'], t2) + \
                        fft_weight * freloss(outputs['img_fre'], t2, mask)


                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

                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 DEBUG:
                        break

                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 = [], [], []

            t2_MSE_first_step, t2_PSNR_first_step, t2_SSIM_first_step = [], [], []

            t1_MSE_krecon, t1_PSNR_krecon, t1_SSIM_krecon = [], [], []
            t2_MSE_krecon, t2_PSNR_krecon, t2_SSIM_krecon = [], [], []
            ids = 0
            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)


                if use_kspace:
                    t = num_timesteps - 1
                    mask = kspace_masks[t]
                    target_fft, _ = apply_tofre(t2.clone(), mask)
                    fft, mask     = apply_tofre(t2_in.clone(), mask)

                    fft = target_fft * mask + fft * (1 - mask)  # Seems too easy
                    t2_in = apply_to_spatial(fft)

                    while t >= 0:
                        if use_time_model:
                            outputs = network(t2_in, t1, t)['img_out']
                        else:
                            outputs = network(t2_in, t1)['img_out']

                        if t == num_timesteps - 1:
                            first_step_recon = outputs

                        if t == 0:
                            mask  = kspace_masks[0]     # last one
                            t2_in = outputs

                        else:
                            k_full = kspace_masks[-1]  # True
                            t2_in_fre, k_full = apply_tofre(t2_in, k_full)

                            with torch.no_grad():

                                kt_sub_1 = kspace_masks[t - 1]  # get_kspace_kernels(t - 2).cuda()
                                kt = kspace_masks[t]            # current one
                                k_residual = kt_sub_1 - kt

                                recon_sample_fre, k_residual = apply_tofre(outputs, k_residual)

                                #  fft = target_fft * mask + fft * (1 - mask)
                                t2_in_fre = t2_in_fre * (1 - k_residual) + recon_sample_fre * k_residual       # substitute

                                outputs = apply_to_spatial(t2_in_fre)
                                t2_in = outputs

                        t = t - 1
                    t2_out = t2_in
                else:
                    t2_out = network(t2_in, t1)['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)
                    t2_first_step_recon_img = (np.clip(first_step_recon.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)
                    t2_first_step_recon_img = (np.clip(first_step_recon.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_first_step_recon_img)
                    PSNR = peak_signal_noise_ratio(t2_img, t2_first_step_recon_img)
                    SSIM = structural_similarity(t2_img, t2_first_step_recon_img)
                    t2_MSE_first_step.append(MSE)
                    t2_PSNR_first_step.append(PSNR)
                    t2_SSIM_first_step.append(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)

                ids += 1
                if ids > 100:
                    break

            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_first_step_mse = np.array(t2_MSE_first_step).mean()
            t2_first_step_psnr = np.array(t2_PSNR_first_step).mean()
            t2_first_step_ssim = np.array(t2_SSIM_first_step).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 First MRI:] average MSE: {t2_first_step_mse} average PSNR: {t2_first_step_psnr} average SSIM: {t2_first_step_ssim}")
            logging.info(f"[T2 MRI:] average MSE: {t2_mse} average PSNR: {t2_psnr} average SSIM: {t2_ssim}")
            print("Snapshot_path = ", snapshot_path)

            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()