import argparse import logging import os import random import sys import numpy as np import torch import torch.backends.cudnn as cudnn from torch.utils.data import DataLoader from torchvision import transforms from networks.unet_model import UNet from dataloaders.dataloader import FundusSegmentation, ProstateSegmentation, MNMSSegmentation, BUSISegmentation import dataloaders.custom_transforms as tr from utils import losses, metrics, ramps, util from medpy.metric import binary from scipy.ndimage import zoom from importlib import import_module from segment_anything import sam_model_registry parser = argparse.ArgumentParser() parser.add_argument('--dataset', type=str, default='prostate', choices=['fundus', 'prostate', 'MNMS', 'BUSI']) parser.add_argument("--save_name", type=str, default="debug", help="experiment_name") parser.add_argument("--overwrite", action='store_true') parser.add_argument("--model", type=str, default="unet", help="model_name") parser.add_argument("--gpu", type=str, default='0') parser.add_argument('--eval',type=bool, default=True) parser.add_argument("--test_bs", type=int, default=1) parser.add_argument('--domain_num', type=int, default=6) parser.add_argument('--lb_domain', type=int, default=1) parser.add_argument('--rank', type=int, default=4, help='Rank for LoRA adaptation') parser.add_argument('--module', type=str, default='sam_lora_image_encoder') parser.add_argument('--img_size', type=int, default=512, help='input patch size of network input') parser.add_argument('--vit_name', type=str, default='vit_b', help='select one vit model') args = parser.parse_args() def to_2d(input_tensor): input_tensor = input_tensor.unsqueeze(1) tensor_list = [] temp_prob = input_tensor == torch.ones_like(input_tensor) tensor_list.append(temp_prob) temp_prob2 = input_tensor > torch.zeros_like(input_tensor) tensor_list.append(temp_prob2) output_tensor = torch.cat(tensor_list, dim=1) return output_tensor.float() def to_3d(input_tensor): input_tensor = input_tensor.unsqueeze(1) tensor_list = [] for i in range(1, 4): temp_prob = input_tensor == i * torch.ones_like(input_tensor) tensor_list.append(temp_prob) output_tensor = torch.cat(tensor_list, dim=1) return output_tensor.float() @torch.no_grad() def test(args, model, test_dataloader, epoch, model_name): model.eval() val_dice = [0.0] * n_part val_dc, val_jc, val_hd, val_asd = [0.0] * n_part, [0.0] * n_part, [0.0] * n_part, [0.0] * n_part domain_num = len(test_dataloader) num = 0 for i in range(domain_num): cur_dataloader = test_dataloader[i] domain_val_dice = [0.0] * n_part domain_val_dc, domain_val_jc, domain_val_hd, domain_val_asd = [0.0] * n_part, [0.0] * n_part, [0.0] * n_part, [0.0] * n_part domain_code = i+1 for batch_num,sample in enumerate(cur_dataloader): assert(domain_code == sample['dc'][0].item()) mask = sample['label'] if args.dataset == 'fundus': lb_mask = (mask<=128) * 2 lb_mask[mask==0] = 1 mask = lb_mask elif args.dataset == 'prostate': mask = mask.eq(0).long() elif args.dataset == 'MNMS': mask = mask.long() elif args.dataset == 'BUSI': mask = mask.eq(255).long() if model_name == 'SAM': data = sample['image'].cuda() output = model(data, multimask_output, args.img_size)['masks'] elif model_name == 'unet': data = sample['unet_size_img'].cuda() output = model(data) pred_label = torch.max(torch.softmax(output,dim=1), dim=1)[1] pred_label = torch.from_numpy(zoom(pred_label.cpu(), (1, patch_size / data.shape[-2], patch_size / data.shape[-1]), order=0)) if args.dataset == 'fundus': pred_label = to_2d(pred_label) mask = to_2d(mask) pred_onehot = pred_label.clone() mask_onehot = mask.clone() elif args.dataset == 'prostate' or args.dataset == 'BUSI': pred_onehot = pred_label.clone().unsqueeze(1) mask_onehot = mask.clone().unsqueeze(1) elif args.dataset == 'MNMS': pred_onehot = to_3d(pred_label) mask_onehot = to_3d(mask) dice = dice_calcu[args.dataset](np.asarray(pred_label.cpu()),mask.cpu()) avg_dice = sum(dice)/len(dice) dc, jc, hd, asd = [0.0] * n_part, [0.0] * n_part, [0.0] * n_part, [0.0] * n_part for j in range(len(data)): for i, p in enumerate(part): dc[i] += binary.dc(np.asarray(pred_onehot[j,i], dtype=bool), np.asarray(mask_onehot[j,i], dtype=bool)) jc[i] += binary.jc(np.asarray(pred_onehot[j,i], dtype=bool), np.asarray(mask_onehot[j,i], dtype=bool)) if pred_onehot[j,i].float().sum() < 1e-4: hd[i] += 100 asd[i] += 100 else: hd[i] += binary.hd95(np.asarray(pred_onehot[j,i], dtype=bool), np.asarray(mask_onehot[j,i], dtype=bool)) asd[i] += binary.asd(np.asarray(pred_onehot[j,i], dtype=bool), np.asarray(mask_onehot[j,i], dtype=bool)) for i, p in enumerate(part): dc[i] /= len(data) jc[i] /= len(data) hd[i] /= len(data) asd[i] /= len(data) for i in range(len(domain_val_dice)): domain_val_dice[i] += dice[i] domain_val_dc[i] += dc[i] domain_val_jc[i] += jc[i] domain_val_hd[i] += hd[i] domain_val_asd[i] += asd[i] for i in range(len(domain_val_dice)): domain_val_dice[i] /= len(cur_dataloader) val_dice[i] += domain_val_dice[i] domain_val_dc[i] /= len(cur_dataloader) val_dc[i] += domain_val_dc[i] domain_val_jc[i] /= len(cur_dataloader) val_jc[i] += domain_val_jc[i] domain_val_hd[i] /= len(cur_dataloader) val_hd[i] += domain_val_hd[i] domain_val_asd[i] /= len(cur_dataloader) val_asd[i] += domain_val_asd[i] text = 'domain%d epoch %d :' % (domain_code, epoch) text += '\n\t' for n, p in enumerate(part): text += 'val_%s_dice: %f, ' % (p, domain_val_dice[n]) text += '\n\t' for n, p in enumerate(part): text += 'val_%s_dc: %f, ' % (p, domain_val_dc[n]) text += '\t' for n, p in enumerate(part): text += 'val_%s_jc: %f, ' % (p, domain_val_jc[n]) text += '\n\t' for n, p in enumerate(part): text += 'val_%s_hd: %f, ' % (p, domain_val_hd[n]) text += '\t' for n, p in enumerate(part): text += 'val_%s_asd: %f, ' % (p, domain_val_asd[n]) logging.info(text) model.train() for i in range(len(val_dice)): val_dice[i] /= domain_num val_dc[i] /= domain_num val_jc[i] /= domain_num val_hd[i] /= domain_num val_asd[i] /= domain_num text = 'epoch %d :' % (epoch) text += '\n\t' for n, p in enumerate(part): text += 'val_%s_dice: %f, ' % (p, val_dice[n]) text += '\n\t' for n, p in enumerate(part): text += 'val_%s_dc: %f, ' % (p, val_dc[n]) text += '\t' for n, p in enumerate(part): text += 'val_%s_jc: %f, ' % (p, val_jc[n]) text += '\n\t' for n, p in enumerate(part): text += 'val_%s_hd: %f, ' % (p, val_hd[n]) text += '\t' for n, p in enumerate(part): text += 'val_%s_asd: %f, ' % (p, val_asd[n]) logging.info(text) return val_dice, val_dc, val_jc, val_hd, val_asd def main(args, snapshot_path): def create_model(model_name=None, ema=False): # Network definition if model_name == 'SAM': logging.info("load from {}".format(args.ckpt)) sam, img_embedding_size = sam_model_registry[args.vit_name](image_size=args.img_size, num_classes=num_classes, checkpoint=args.ckpt, pixel_mean=[0, 0, 0], pixel_std=[1, 1, 1]) pkg = import_module(args.module) model = pkg.LoRA_Sam(sam, args.rank) if ema: for param in model.parameters(): param.detach_() return model.cuda(), img_embedding_size elif model_name == 'unet': model = UNet(n_channels = num_channels, n_classes = num_classes+1) if ema: for param in model.parameters(): param.detach_() return model.cuda() else: raise Exception('Please provide model name.') model, img_embedding_size = create_model(model_name='SAM') low_res = img_embedding_size * 4 normal_toTensor = transforms.Compose([ tr.Normalize_tf(dataRange=[0,1]), tr.ToTensor(low_res=low_res, unet_size=patch_size) ]) domain_num = args.domain_num test_dataset = [] test_dataloader = [] for i in range(1, domain_num+1): cur_dataset = dataset(base_dir=train_data_path, phase='test', splitid=-1, domain=[i], normal_toTensor=normal_toTensor, img_size=args.img_size) test_dataset.append(cur_dataset) for i in range(0,domain_num): cur_dataloader = DataLoader(test_dataset[i], batch_size = args.test_bs, shuffle=False, num_workers=0, pin_memory=True) test_dataloader.append(cur_dataloader) if args.eval: args.lb_domain = 1 if args.model == 'unet': model = create_model(model_name='unet') model.load_state_dict(torch.load('../model/{}/{}/unet_avg_dice_best_model.pth'.format(args.dataset, args.save_name))) elif args.model == 'SAM': model, img_embedding_size = create_model(model_name='SAM') model.load_lora_parameters('../model/{}/{}/SAM_avg_dice_best_model.pth'.format(args.dataset, args.save_name)) test(args, model,test_dataloader,args.lb_domain, args.model) exit() if __name__ == "__main__": snapshot_path = "../model/" + args.dataset + "/" + args.save_name + "/" if args.dataset == 'fundus': num_channels = 3 num_classes = 2 train_data_path='../data/Fundus' patch_size = 256 part = ['cup', 'disc'] dataset = FundusSegmentation if args.domain_num >=4: args.domain_num = 4 elif args.dataset == 'prostate': num_channels = 1 num_classes = 1 train_data_path="../data/ProstateSlice" patch_size = 384 part = ['base'] dataset = ProstateSegmentation if args.domain_num >= 6: args.domain_num = 6 elif args.dataset == 'MNMS': num_channels = 1 num_classes = 3 train_data_path="../data/mnms" patch_size = 288 part = ['lv', 'myo', 'rv'] dataset = MNMSSegmentation if args.domain_num >= 4: args.domain_num = 4 elif args.dataset == 'BUSI': num_channels = 1 num_classes = 1 train_data_path="../data/Dataset_BUSI_with_GT" patch_size = 256 part = ['base'] dataset = BUSISegmentation if args.domain_num >= 2: args.domain_num = 2 if num_classes > 1: multimask_output = True else: multimask_output = False n_part = len(part) dice_calcu = {'fundus':metrics.dice_coeff_2label, 'prostate':metrics.dice_coeff, 'MNMS':metrics.dice_coeff_3label, 'BUSI':metrics.dice_coeff} args.ckpt = '../checkpoints/medsam_vit_b.pth' os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu if not os.path.exists(snapshot_path): os.makedirs(snapshot_path) elif not args.overwrite: raise Exception('file {} is exist!'.format(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)) cmd = " ".join(["python"] + sys.argv) logging.info(cmd) logging.info(str(args)) main(args, snapshot_path)