test_diff.py

import os

import numpy
import numpy as np
import torchvision
import torch
import time
import argparse
import torchvision.transforms as transforms
from torch.optim import Adam
from models.denoiseModel.UNet import *
from models.diffusionModel.Diffusion import *
from utils.trainer import *
from matplotlib import pyplot as plt
from dataloaders.datasets import *
from visdom import Visdom
from torch.utils.data import DataLoader


def get_parser():
    parser = argparse.ArgumentParser()

    parser.add_argument("--batch_size", type=int, default=1, help='Test batch size')
    parser.add_argument("--channels", type=int, default=1, help='Number of channels')
    parser.add_argument("--img_size", type=int, default=240, help='Image size (will be resized to this size)')
    parser.add_argument("--dax", type=str, default="BraTS18", choices=['Abdominal', 'BraTS18', 'Fundus'], help='Dataset name')
    parser.add_argument("--data_dir", type=str,
                        default="./datasets/BraTS_2018", help='Dataset directory')
    parser.add_argument("--root_path", type=str, default="./saved_train_models/diff_model/brats2018/t2/model.pth",
                        help='Diffusion model checkpoint path')
    parser.add_argument("--source_domain", type=str, default="t2", help='Source domain')

    parser.add_argument("--data_transform", help='Whether to preprocess the data', action='store_true')

    parser.add_argument("--denoise_mode", type=str, default="UNetFree", help='Denoising model type')
    parser.add_argument("--dim_mults", type=tuple, default=(1, 2, 4, 8), help='Channel multipliers for each model stage')
    parser.add_argument("--dim_bn", type=int, default=32, help='Base channel dimension')
    parser.add_argument("--betas_schedule_name", type=str, default="linear", help='Noise schedule')
    parser.add_argument("--time_steps", type=int, default=1000, help='Number of diffusion timesteps')
    parser.add_argument("--beta_start", type=float, default=0.0001, help='Starting beta (noise variance)')
    parser.add_argument("--beta_end", type=float, default=0.02, help='Ending beta (noise variance)')
    parser.add_argument("--classifier_mode", type=str, default="free", help='Conditional guidance mode')

    parser.add_argument("--epochs", type=int, default=100, help='Number of training epochs')
    parser.add_argument("--lr_diff", type=float, default=1e-4, help='Denoising model learning rate')

    parser.add_argument("--visdom_flg", help='Enable visualization via Visdom', action='store_true')
    parser.add_argument("--visdom_port", type=int, default=8097, help='Visdom port')

    parser.add_argument("--allsave_path", type=str, default="./log/allsave.txt", help='Log file for args and corresponding paths')
    parser.add_argument("--bz_mode", type=str, default="none", choices=['none', '1', '2', '3', '4'],
                        help='Bezier transform mode')

    parser.add_argument("--plt_mode", action='store_true', help='Whether to plot')

    parser.add_argument('--gpu', type=str, default='0', help='GPU id')

    args = parser.parse_args()
    show_args(args)
    return args


def show_args(args):
    for (key, value) in args.__dict__.items():
        print(key, value)


def get_cmd(args):
    strcmd = 'python test_diff.py'
    for (key, value) in args.__dict__.items():
        if value is False:
            continue
        if value is True:
            strcmd += f' --{key}'
        else:
            strcmd += f' --{key} {value}'
    return strcmd

def min_max_normalize_tensor(tensor):
    min_val = torch.min(tensor)
    max_val = torch.max(tensor)
    if max_val - min_val == 0:
        return tensor
    normalized_tensor = (tensor - min_val) / (max_val - min_val)
    return normalized_tensor

def run(args):
    # cuda
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
    device = "cuda" if torch.cuda.is_available() else "cpu"

    if not os.path.exists(args.data_dir):
        os.makedirs(args.data_dir)

    tran_list = None
    tran_train = None
    if args.data_transform:
        tran_list = [transforms.Resize((args.img_size, args.img_size)), transforms.ToTensor()]
        tran_train = transforms.Compose(tran_list)

    if args.dax == 'BraTS18':
        dataset = BraTS18Dataset(args.data_dir, dana=get_dom(args=args))
    elif args.dax == 'Fundus':
        dataset = FundusDataset(args.data_dir, dana=get_dom(args=args), img_size=args.img_size)
    else:
        raise NotImplementedError
    data_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=0)

    # # denoise model
    denoise_model = None
    if args.denoise_mode == 'UNetFree':
        denoise_model = UNetFree(in_dim=args.dim_bn, dim_mults=args.dim_mults, channels=args.channels,
                                 cond_channel=args.channels).to(device)
    else:
        raise NotImplementedError
    # # diffusion tool
    diffusion = Diffusion(shape=(args.channels, args.img_size, args.img_size), denoise_model=denoise_model,
                          betas_schedule_name=args.betas_schedule_name, time_steps=args.time_steps,
                          beta_start=args.beta_start, \
                          beta_end=args.beta_end, classifier_mode=args.classifier_mode)

    diffusion.load_state_dict(torch.load(args.root_path),strict=False)
    diffusion.eval()
    with torch.no_grad():
        for i, (img, seg, _) in enumerate(data_loader):
            if args.bz_mode != 'none':
                fe_tmp = np.array(img.cpu())
                fe_tmp = fe_tmp.transpose((2, 3, 0, 1))
                fe_tmp = fe_tmp.squeeze()

                mask_tmp = np.array(seg.cpu(), dtype=np.int64)

                mask_tmp = mask_tmp.transpose((1, 2, 0))
                mask_tmp = mask_tmp.squeeze()
                if len(fe_tmp.shape) == 2:
                    fe_tmp = np.expand_dims(fe_tmp, axis=-1)
                    mask_tmp = np.expand_dims(mask_tmp, axis=-1)

                if args.bz_mode == '1':
                    fe_tmp = get_bezier_img(fe_tmp, mask_tmp)
                elif args.bz_mode == '2':
                    fe_tmp = get_bezier_img2(fe_tmp, mask_tmp)
                elif args.bz_mode == '3':
                    fe_tmp = get_bezier_img3(fe_tmp, mask_tmp)
                elif args.bz_mode == '4':
                    fe_tmp = get_bezier_img4(fe_tmp, mask_tmp)
                fe_tmp = np.expand_dims(fe_tmp, axis=-1)
                fe_tmp = fe_tmp.transpose((2, 3, 0, 1))
                for i in range(img.shape[0]):
                    plt.subplot(121), plt.imshow(fe_tmp[i, 0, :, :], cmap='gray')
                    plt.subplot(122), plt.imshow(img[i, 0, :, :].cpu().numpy(), cmap='gray')
                    plt.axis('off')
                    plt.show()
                img = torch.from_numpy(fe_tmp).float()

            img = img.to(device)
            features, cond = to_fft(img)

            fx = torch.log(features)
            fx2 = torch.exp(fx)

            cond = cond.float()
            samples = diffusion.ddim_sample_loop(nums=img.shape[0], sstep=200, cond=cond)
            gn = samples[-1]
            gn = torch.tensor(gn, dtype=torch.float32)
            gn = gn.to(device)

            gi = to_ifft(torch.exp(gn), cond)
            gi = torch.real(gi)
            gi = normalize_tensor(gi)
            gi = np.array(gi.cpu())
            gi = np.real(gi)

            for i in range(img.shape[0]):
                show_img = gi[i]
                print(show_img.shape)
                if args.dax == 'BraTS18':
                    plt.subplot(321), plt.imshow(show_img[i], cmap="gray"), plt.title('ZH')
                    plt.subplot(322), plt.imshow(img[i].cpu().squeeze(), cmap="gray"), plt.title('YT')
                    plt.subplot(323), plt.imshow(gn[i].cpu().squeeze(), cmap="gray"), plt.title('LGE')
                    plt.subplot(324), plt.imshow(torch.log(features[i]).cpu().squeeze(), cmap="gray"), plt.title('LFT')
                    plt.subplot(325), plt.imshow(torch.exp(gn[i]).cpu().squeeze(), cmap="gray"), plt.title('GE')
                    plt.subplot(326), plt.imshow(features[i].cpu().squeeze(), cmap="gray"), plt.title('FT')

                plt.axis('off')
                plt.show()


if __name__ == '__main__':
    run(get_parser())