test_codg.py

import copy
import os
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
from models.segModel.Net_Factory import *
from models.segModel.DNASeg import *
from models.segModel.FFTSeg import *
import medpy.metric.binary as mmb


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

    parser.add_argument('--test_domain_list', nargs='+', type=str)

    parser.add_argument("--batch_size", type=int, default=16, 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", help='Dataset name')
    parser.add_argument("--data_dir", type=str, default="./datasets/BraTS_2018", help='Dataset directory')
    parser.add_argument("--data_transform", help='Whether to preprocess the data', action='store_true')

    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 for each model stage')
    parser.add_argument("--seg_path", type=str, default="./saved_train_models/dna_model/2/model_15.pth",
                        help='Model checkpoint path')
    parser.add_argument("--seg_model", type=str, default="SegUNet", help='Model name')
    parser.add_argument("--denoise_mode", type=str, default="UNetFree", help='Denoising model type')
    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("--noise_mu", type=float, default=1.0, help='Noise scale / noise base factor')
    parser.add_argument("--sstep", type=int, default=100, help='Number of DDIM sampling steps')
    parser.add_argument("--rand_limit", type=float, default=0, help='Probability of using original data (0-1)')
    parser.add_argument("--bn_flg", help='Whether to normalize the data', action='store_true')
    parser.add_argument("--seg_bn", type=int, default=32, help='Base channel dimension of the segmentation model')
    parser.add_argument("--dna_flg", help='Whether to enable DNA double-strand training mode (enabled by default)',
                        action='store_true')

    parser.add_argument("--out_channel", type=int, default=2,
                        help='Number of output channels of the segmentation model')
    parser.add_argument("--restore_out_channel", type=int, default=1,
                        help='Number of output channels of the restoration model')

    parser.add_argument("--look_img", help='Whether to visualize comparison results', action='store_true')
    parser.add_argument("--diff_flg", help='Whether to perform diffusion before evaluation', action='store_true')

    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("--tqdm_flg", help='Enable tqdm progress bar', action='store_true')

    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 cls_test.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 args_on_visdom(vis, args):
    vis.text(f'args', win='args_text')
    for (key, value) in args.__dict__.items():
        vis.text(f'{key}: {value}', win='args_text', append=True)
    vis.text(f'==============================', win='args_text', append=True)
    vis.text(f'start by:', win='args_text', append=True)
    vis.text(f'{get_cmd(args)}', win='args_text', append=True)


def write_args(id_path, args):
    with open(args.allsave_path, 'a+') as fi:
        fi.write(f'{id_path}\n')
        fi.write('------------------------------\n')
        for (key, value) in args.__dict__.items():
            fi.write(f'{key}: {value}\n')
        fi.write('------------------------------\n')
        fi.write(get_cmd(args))
        fi.write('\n')
        fi.write('------------------------------\n')
        fi.write(args.root_path + '/' + str(id_path))
        fi.write('\n')
        fi.write('==============================\n')


def get_all_dice(pred, gt, args):
    dice_all = []
    for i in range(args.out_channel):
        tmp_p = np.array(pred == i, dtype=np.int64)
        tmp_y = np.array(gt == i, dtype=np.int64)
        dice_all.append(mmb.dc(tmp_p, tmp_y))
    return dice_all


def get_all_hd95(pred, gt, args):
    hd95_all = []
    for i in range(args.out_channel):
        tmp_p = np.array(pred == i, dtype=np.int64)
        tmp_y = np.array(gt == i, dtype=np.int64)
        if tmp_p.sum() == 0 or tmp_y.sum() == 0:
            hd95_all.append(100)
        else:
            hd95_all.append(mmb.hd95(tmp_p, tmp_y))
    return hd95_all


def run(args, model_path):
    cond_name = ['CondSegNet', 'CondMulSegNet', 'CondMulSegNet2', 'CondSegNet2']
    # cuda
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
    device = "cuda" if torch.cuda.is_available() else "cpu"

    test_domain_list = args.test_domain_list
    num_domain = len(test_domain_list)

    for test_idx in range(num_domain):

        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, mode='test', dana=[test_domain_list[test_idx]])
        elif args.dax == 'Fundus':
            dataset = FundusDataset(args.data_dir, mode='test', dana=[test_domain_list[test_idx]], 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, rand_mu=args.noise_mu)

        seg_model = net_factory(args=args, device=device)
        modified_args = copy.deepcopy(args)
        modified_args.out_channel = args.restore_out_channel
        restore_model = net_factory(args=modified_args, device=device)

        fftseg = DNASeg(seg_model=seg_model, restore_model=restore_model, diffusion=diffusion, sstep=args.sstep, rand_limit=args.rand_limit)
        fftseg.load_state_dict(torch.load(model_path))

        fftseg.eval()
        all_pr_y = None
        all_y = None
        with torch.no_grad():
            if args.tqdm_flg:
                loop = tqdm(enumerate(data_loader), total=len(data_loader))
            else:
                loop = enumerate(data_loader)
            for step, (x, y) in loop:
                x = x.to(device)
                x = x.float()
                if args.diff_flg:
                    x_recon = fftseg.get_recon_img(x)
                    x_recon = torch.tensor(x_recon, dtype=torch.float32)
                    x_recon = x_recon.to(x.device)
                    x = x_recon
                if args.bn_flg:
                    x = normalize_tensor(x)
                # print(y.max())
                if args.seg_model in cond_name:
                    px, mx = to_fft(x)
                    px = torch.log(px)
                    pre_y, _ = fftseg.seg_model(x, px)
                else:
                    pre_y, _ = fftseg.seg_model(x)

                y = y.squeeze(1).numpy()
                pre_y = pre_y.cpu().detach().numpy()
                pre_y = np.argmax(pre_y, axis=1)

                if all_pr_y is None:
                    all_pr_y = pre_y
                    all_y = y
                else:
                    all_pr_y = np.concatenate((all_pr_y, pre_y), axis=0)
                    all_y = np.concatenate((all_y, y), axis=0)

                if args.tqdm_flg:
                    loop.set_description(f'[{step}/{len(data_loader)}]')


        dice_all = get_all_dice(all_pr_y, all_y, args=args)
        hd95_all = get_all_hd95(all_pr_y, all_y, args=args)
        sum_dc = 0
        sum_hd95 = 0
        for i in range(args.out_channel):
            print('dc:', dice_all[i], 'hd95', hd95_all[i])
            if i > 0:
                sum_dc += dice_all[i]
                sum_hd95 += hd95_all[i]
        print('model-name', model_path, 'domain:', test_domain_list[test_idx], 'avg_dc', sum_dc / (args.out_channel - 1), 'avg_hd95', sum_hd95 / (args.out_channel - 1))

if __name__ == '__main__':
    args = get_parser()
    model_path = args.seg_path
    run(args, model_path)