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

import os
import sys
from tqdm import tqdm
from tensorboardX import SummaryWriter

import logging
import time
import torch.optim as optim
from torch.utils.data import DataLoader
from networks.mynet import TwoBranch
from networks_time.mynet import DiffTwoBranch

from utils.option import args

from dataloaders.fastmri import build_dataset
from frequency_diffusion.degradation.k_degradation import get_ksu_kernel, apply_tofre, apply_to_spatial
from utils.lpips import LPIPS
from utils.metric import nmse, psnr, ssim, AverageMeter
from collections import defaultdict

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
from utils.utils import *

# --num_timesteps 30 --image_size 320 --use_kspace True --use_time_model True --ACCELERATIONS 4 --gpu 0 --phase train
DEBUG = False
use_kspace = args.use_kspace
frequency_distortion = False
use_time_model = args.use_time_model
num_timesteps = args.num_timesteps
image_size = 320
distortion_sigma = 10 / 255

if args.phase == 'test':
    kspace_masks = np.load(f"./dataloaders/example_mask/kspace_{args.ACCELERATIONS[0]}_mask.npy")
    kspace_masks = torch.from_numpy(np.asarray(kspace_masks)).cuda()

else:
    # Output a list of k-space kernels
    kspace_masks = get_ksu_kernel(num_timesteps, image_size,
                                  ksu_routine="LogSamplingRate",
                                  accelerated_factor=args.ACCELERATIONS[0],
                                  )  # args.ACCELERATIONS = [4] or [8]

    np.save(f"./dataloaders/example_mask/kspace_{args.ACCELERATIONS[0]}_mask.npy", kspace_masks)
kspace_masks = torch.from_numpy(np.asarray(kspace_masks)).cuda()

print("kspace kernels shape:", kspace_masks.shape)  # (1, 1, 320, 320)


@torch.no_grad()
def evaluate(model, data_loader, device):
    model.eval()

    nmse_meter, psnr_meter, ssim_meter = AverageMeter(), AverageMeter(), AverageMeter()
    direct_nmse, direct_psnr, direct_ssim = AverageMeter(), AverageMeter(), AverageMeter()
    output_dic = defaultdict(dict)
    target_dic = defaultdict(dict)
    input_dic = defaultdict(dict)
    direct_dic = defaultdict(dict)

    for id, data in enumerate(data_loader):
        pd, pdfs, _ = data
        name = os.path.basename(pdfs[4][0]).split('.')[0]

        target = pdfs[1].to(device)
        mean, std = pdfs[2], pdfs[3]

        fname = pdfs[4]
        slice_num = pdfs[5]

        mean = mean.unsqueeze(1).unsqueeze(2).to(device)
        std = std.unsqueeze(1).unsqueeze(2).to(device)

        pd_img = pd[1].unsqueeze(1).to(device)
        pdfs_img = pdfs[0].unsqueeze(1).to(device)

        pdfs_img_origin = pdfs_img.clone()

        # Degradation
        if use_kspace:
            b = pd_img.size(0)
            t = torch.randint(num_timesteps - 1, num_timesteps, (b,), device=device).long()  # t-1
            mask = kspace_masks[t]
            fft, mask = apply_tofre(target.clone(), mask)
            fft = fft * mask + 0.0
            pdfs_img = apply_to_spatial(fft)
            # print("pdfs_img shape:", pdfs_img.shape, pdfs_img.min(), pdfs_img.max())

            while t >= 0:
                if use_time_model:
                    outputs = model(pdfs_img, pd_img, t)['img_out']
                else:
                    outputs = model(pdfs_img, pd_img)['img_out']

                if t == num_timesteps - 1:
                    direct_recon = outputs

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

                else:
                    k_full = kspace_masks[-1]
                    faded_recon_sample_fre, k_full = apply_tofre(pdfs_img, k_full)

                    with torch.no_grad():

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

                        recon_sample_fre, k_residual = apply_tofre(outputs, k_residual)
                        fre_amend = recon_sample_fre * k_residual
                        faded_recon_sample_fre = faded_recon_sample_fre + fre_amend  # * (1-k_residual)
                        # faded_recon_sample_fre = faded_recon_sample_fre * kt_sub_1
                        outputs = apply_to_spatial(faded_recon_sample_fre)
                        pdfs_img = outputs

                t = t - 1

        else:
            outputs = model(pdfs_img, pd_img)['img_out']

        target = target * std + mean
        inputs = pdfs_img_origin.squeeze(1) * std + mean
        outputs = outputs.squeeze(1) * std + mean
        direct_recon = direct_recon.squeeze(1) * std + mean

        # print("target/outputs shape: ", target.shape, outputs.shape)
        # our_nmse = nmse(target[0].cpu().numpy(), outputs[0].cpu().numpy())
        # our_psnr = psnr(target[0].cpu().numpy(), outputs[0].cpu().numpy())
        # our_ssim = ssim(target[0].cpu().numpy(), outputs[0].cpu().numpy())

        # print('name:{}, slice:{}, nmse:{}, psnr:{}, ssim:{}'.format(name, slice_num[0], our_nmse, our_psnr, our_ssim))

        for i, f in enumerate(fname):
            output_dic[f][slice_num[i]] = outputs[i]
            target_dic[f][slice_num[i]] = target[i]
            input_dic[f][slice_num[i]] = inputs[i]
            direct_dic[f][slice_num[i]] = direct_recon[i]

        if id > 100:
            break

    for name in output_dic.keys():
        f_output = torch.stack([v for _, v in output_dic[name].items()])
        f_target = torch.stack([v for _, v in target_dic[name].items()])
        our_nmse = nmse(f_target.cpu().numpy(), f_output.cpu().numpy())
        our_psnr = psnr(f_target.cpu().numpy(), f_output.cpu().numpy())
        our_ssim = ssim(f_target.cpu().numpy(), f_output.cpu().numpy())

        nmse_meter.update(our_nmse, 1)
        psnr_meter.update(our_psnr, 1)
        ssim_meter.update(our_ssim, 1)

        direct_nmse.update(
            nmse(f_target.cpu().numpy(), torch.stack([v for _, v in direct_dic[name].items()]).cpu().numpy()), 1)
        direct_psnr.update(
            psnr(f_target.cpu().numpy(), torch.stack([v for _, v in direct_dic[name].items()]).cpu().numpy()), 1)
        direct_ssim.update(
            ssim(f_target.cpu().numpy(), torch.stack([v for _, v in direct_dic[name].items()]).cpu().numpy()), 1)

    print("==> Evaluate Metric")
    print("Direct Results ----------")
    print("NMSE: {:.4}".format(direct_nmse.avg))
    print("PSNR: {:.4}".format(direct_psnr.avg))
    print("SSIM: {:.4}".format(direct_ssim.avg))
    print("------------------")

    print("==> Evaluate Metric")
    print("Results ----------")
    print("NMSE: {:.4}".format(nmse_meter.avg))
    print("PSNR: {:.4}".format(psnr_meter.avg))
    print("SSIM: {:.4}".format(ssim_meter.avg))
    print("------------------")
    model.train()

    return {'NMSE': nmse_meter.avg, 'PSNR': psnr_meter.avg, 'SSIM': ssim_meter.avg}


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 frequency_distortion:
        snapshot_path   = snapshot_path.rstrip("/") + '_no_distortion/'

    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()
    # network = build_model_from_name(args).cuda()
    device = torch.device('cuda')
    network.to(device)
    lpips_loss = LPIPS().eval().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 = build_dataset(args, mode='train', use_kspace=use_kspace)
    db_test = build_dataset(args, mode='val', use_kspace=use_kspace)

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

    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 = {'NMSE': 10000000, 'PSNR': 0, 'SSIM': 0}
        fft_weight = 0.01
        criterion = nn.L1Loss().to(device, non_blocking=True)
        freloss = Frequency_Loss().to(device, non_blocking=True)

        for epoch_num in tqdm(range(max_epoch), ncols=70):
            time1 = time.time()
            start_time = time.time()
            for i_batch, sampled_batch in enumerate(trainloader):
                time2 = time.time()

                pd, pdfs, _ = sampled_batch
                target = pdfs[1]

                mean, std = pdfs[2], pdfs[3]

                pd_img = pd[1].unsqueeze(1)
                pdfs_img = pdfs[0].unsqueeze(1)
                target = target.unsqueeze(1)

                b = pd_img.size(0)

                pd_img = pd_img.to(device)  # [4, 1, 320, 320]
                pdfs_img = pdfs_img.to(device)  # [4, 1, 320, 320]
                target = target.to(device)  # [4, 1, 320, 320]

                time3 = time.time()

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

                    fft, mask = apply_tofre(target.clone(), mask)
                    fft = fft * mask

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

                    pdfs_img = apply_to_spatial(fft)

                # breakpoint()
                if use_time_model:
                    outputs = network(pdfs_img, pd_img, t)
                else:
                    outputs = network(pdfs_img, pd_img)

                loss = criterion(outputs['img_out'], target) + \
                       fft_weight * freloss(outputs['img_fre'], target, mask) + \
                       criterion(outputs['img_fre'], target) + \
                       0.01 * lpips_loss(outputs['img_out'], target).mean()

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

                time5 = time.time()

                # summary
                iter_num = iter_num + 1

                print_iter = 100  # if not DEBUG else 5
                if iter_num % print_iter == 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()
            eval_result = evaluate(network, testloader, device)

            if eval_result['PSNR'] > best_status['PSNR']:
                best_status = {'NMSE': eval_result['NMSE'], 'PSNR': eval_result['PSNR'], 'SSIM': eval_result['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 saved:', best_checkpoint_path)

            logging.info(
                f"average MSE: {eval_result['NMSE']} average PSNR: {eval_result['PSNR']} average SSIM: {eval_result['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()