MRI_recon/code/Frequency-Diffusion/FSMNet/train_m4raw.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
import matplotlib.pyplot as plt

use_new_dataloader = True

if use_new_dataloader:
    from dataloaders.m4raw_std_dataloader import M4Raw_TestSet, M4Raw_TrainSet, normalize, normalize_instance_dim
else:
    from dataloaders.m4raw_dataloader import M4Raw_TestSet, M4Raw_TrainSet, normalize, normalize_instance_dim

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
# import imsave
from skimage.io import imsave

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 *

frequency_distortion = True



num_timesteps = args.num_timesteps
image_size    = args.image_size
distortion_sigma = 10/255
DEBUG = args.DEBUG
use_kspace = args.use_kspace
use_time_model = args.use_time_model


# Baseline ------------------------------------
# NMSE: 3.3329 ± 0.4915
# PSNR: 34.0418 ± 0.6790
# SSIM: 0.8699 ± 0.0143
# ------------------------------------
# Save Path: model/FSMNet_m4raw_4x//result_case/

# ------------------------------------
# NMSE: 3.1966 ± 0.3939
# PSNR: 34.2115 ± 0.5978
# SSIM: 0.8910 ± 0.0118
# ------------------------------------
# Save Path: model/FSMNet_m4raw_4x_kspace//result_case/


# num_timesteps = 5
image_size    = 240

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

use_in_mean_std = False

# 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/m4raw_{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()
    print_i = 1

    nmse_meter = AverageMeter()
    psnr_meter = AverageMeter()
    ssim_meter = AverageMeter()
    output_dic = defaultdict(dict)
    target_dic = defaultdict(dict)
    input_dic = defaultdict(dict)

    for id, sampled_batch in enumerate(data_loader):

        if use_new_dataloader:
            t1_img, t1_in = sampled_batch['t1'], sampled_batch['t1_in']
            t2_img, t2_in = sampled_batch['t2'], sampled_batch['t2_in']
        else:
            t1_img, t1_in = sampled_batch['ref_image_full'], sampled_batch['ref_image_sub']
            t2_img, t2_in = sampled_batch['tag_image_full'], sampled_batch['tag_image_sub']


        t1_img = t1_img.to(device)
        # t1_in = t1_in.to(device)
        t2_img = t2_img.to(device)
        t2_in = t2_in.to(device)

        mean, std = sampled_batch['t2_mean'], sampled_batch['t2_std']

        fname = sampled_batch['fname']
        slice_num = sampled_batch['slice']

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

        t2_in_origin = t2_in.clone()

        # Degradation
        if use_kspace:
            b = 1
            t = torch.randint(num_timesteps - 1, num_timesteps, (b,), device=device).long()  # t-1
            mask = kspace_masks[t]
            fft, mask = apply_tofre(t2_in.clone(), mask)  # t2_img
            fft = fft * mask + 0.0
            t2_in = apply_to_spatial(fft)

            # Save for debug:
            # put t2_in and t2_in_oringin side by side


            t2_in_origin = t2_in.clone()

            # print("test fft/mask = ", fft.shape, fft.shape)


            if use_in_mean_std:
                t2_in = t2_in * std + mean
                t2_img = t2_img * std + mean
                # print("Test After restore:", t2_in.shape, t2_img.shape)

                t2_in, mean, std = normalize_instance_dim(t2_in, eps=1e-11)
                t2_img  = normalize(t2_img, mean=mean, stddev=std, eps=1e-11)
                t2_in = t2_in.float()
                t2_img = t2_img.float()

                mean = mean[0]
                std = std[0]

                # print("Test Re-normalize:", t2_in.shape, t2_img.shape)

            # print("in put  t2_in shape:", t2_in.shape, t2_in_origin.shape)

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

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

                else:
                    k_full = kspace_masks[-1]
                    faded_recon_sample_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]             #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)
                        t2_in = outputs

                t = t-1

        else:
            outputs = model(t2_in, t1_img)['img_out']

        if print_i:
            t2_in_save = torch.cat([t2_in, t2_in_origin, t2_img], dim=3).cpu().numpy()[0, 0]

            t2_in_save = (t2_in_save - t2_in_save.min()) / (t2_in_save.max() - t2_in_save.min())
            # t2_in_save = np.stack([t2_in_save, t2_in_save, t2_in_save], axis=2)
            # save to file
            os.makedirs("./debug", exist_ok=True)
            save_path = f"./debug/{use_kspace}_{fname[0]}_{slice_num[0]}.png"
            plt.imsave(save_path, t2_in_save, cmap='gray')
            print_i = 0
            print("print_i")


        t2_img  = t2_img.squeeze(1) * std + mean
        inputs  = t2_in_origin.squeeze(1) * std + mean
        outputs = outputs.squeeze(1) * std + mean

        # print("output:", t2_img.shape, outputs.shape, inputs.shape)

        for i, f in enumerate(fname):

            output_dic[f][slice_num[i]] = outputs[i]
            target_dic[f][slice_num[i]] = t2_img[i]
            input_dic[f][slice_num[i]]  = inputs[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)

    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:
        if use_new_dataloader:
            snapshot_path = snapshot_path.rstrip("/") + f'_t{num_timesteps}_new_kspace/'
        else:
            snapshot_path = snapshot_path.rstrip("/") + f'_t{num_timesteps}/'


    if not isinstance(args.test_tag, type(None)):
        snapshot_path = snapshot_path.rstrip("/") + f'_{args.test_tag}/'

    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:
        model = DiffTwoBranch(args).cuda()
    else:
        model = TwoBranch(args).cuda()

    # model = build_model_from_name(args).cuda()
    device = torch.device('cuda')
    model.to(device)
    lpips_loss = LPIPS().eval().to(device)

    if len(args.gpu.split(',')) > 1:
        model = nn.DataParallel(model)
        # model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
    
    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print('number of params: %.2f M' % (n_parameters / 1024 / 1024))

    # db_train = M4Raw_TrainSet(args.root_path, args.MRIDOWN, use_kspace=use_kspace) # build_dataset(args, mode='train')
    # db_test  = M4Raw_TestSet(args.root_path, args.MRIDOWN, use_kspace=use_kspace) # build_dataset(args, mode='val')

    db_train = M4Raw_TrainSet(args, use_kspace=use_kspace, DEBUG=DEBUG)  # build_dataset(args, mode='train')
    db_test  = M4Raw_TestSet(args, use_kspace=use_kspace, DEBUG=DEBUG)  #

    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)

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

        params = list(model.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)
        t = 0

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

                # T1 is the reference image, T2 is the target image
                t1_img, t1_in = sampled_batch['t1'], sampled_batch['t1_in']
                t2_img, t2_in = sampled_batch['t2'], sampled_batch['t2_in']

                t1_img = t1_img.to(device)
                t1_in  = t1_in.to(device)
                t2_img = t2_img.to(device)
                t2_in  = t2_in.to(device)

                time3 = time.time()

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

                    fft, mask = apply_tofre(t2_in.clone(), mask)       # TODO t2_img
                    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)

                    t2_in = apply_to_spatial(fft)

                    if use_in_mean_std:
                        mean, std = sampled_batch['t2_mean'], sampled_batch['t2_std']
                        mean = mean.unsqueeze(1).unsqueeze(1).to(device)
                        std = std.unsqueeze(1).unsqueeze(1).to(device)

                        t2_in = t2_in * std + mean
                        t2_img = t2_img * std + mean
                        # print("Train After restore:", t2_in.shape, t2_img.shape, mean.shape, std.shape)

                        t2_in, mean, std = normalize_instance_dim(t2_in, eps=1e-11)
                        t2_img = normalize(t2_img, mean=mean, stddev=std, eps=1e-11).detach()
                        t2_in = t2_in.float()
                        t2_img = t2_img.float()
                        # print("Train After restore:", t2_in.shape, t2_img.shape)

                # breakpoint()
                if use_time_model:
                    outputs = model(t2_in, t1_img, t) # ['img_out']
                else:
                    outputs = model(t2_in, t1_img)    # ['img_out']

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

                        # 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_(model.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': model.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(model, 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': model.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': model.state_dict()},
                   save_mode_path)
        logging.info("save model to {}".format(save_mode_path))
        writer.close()