def main(args):
if not os.path.exists(args.save_path):
os.mkdir(args.save_path)
net = UNet(n_channels=3, n_classes=1)
checkpoint = flow.load(args.pretrained_path)
net.load_state_dict(checkpoint)
net.to("cuda")
x_test_dir, y_test_dir = get_datadir_path(args, split="test")
test_dataset = Dataset(
x_test_dir, y_test_dir, augmentation=get_test_augmentation(),
)
print("Begin Testing...")
for i, (image, mask) in enumerate(tqdm(test_dataset)):
show_image = image
with flow.no_grad():
image = image / 255.0
image = image.astype(np.float32)
image = flow.tensor(image, dtype=flow.float32)
image = image.permute(2, 0, 1)
image = image.to("cuda")
pred = net(image.unsqueeze(0).to("cuda"))
pred = pred.numpy()
pred = pred > 0.5
save_picture_name = os.path.join(args.save_path, "test_image_" + str(i))
visualize(
save_picture_name, image=show_image, GT=mask[0, :, :], Pred=pred[0, 0, :, :]
)
def run(criterion, net_name, use_dataset_index, use_mask, save_img_pred: bool):
path_data = config.path_data[use_dataset_index]
# 文件路径
path_data_root = path_data["dataset"]
path_checkpoints = path_data["checkpoints"]
path_result = os.path.join(path_data["result"], net_name)
os.makedirs(path_result, exist_ok=True)
net = UNet().to(config.device)
for param in net.parameters():
param.requires_grad = False
path_pert = "/home/pengzx/deepLearning/result/Glaucoma/UNet/pert/"
iou_total = []
for index in config.checkpoints_indexes:
path_pert_save = os.path.join(path_pert, str(index))
attacked_dataset = MyAttackedDataset(path_data_root=path_data_root,
phase="train",
path_pert=path_pert_save,
transform_list=transform_compose)
attacked_data_loader = DataLoader(attacked_dataset,
batch_size=1,
shuffle=False)
for index_2 in config.checkpoints_indexes:
net.load_state_dict(
torch.load(os.path.join(path_checkpoints,
"{}_{}.pth").format(net_name, index_2),
map_location=config.device))
iou_list = []
for i, (img, label, mask, pert,
name) in enumerate(attacked_data_loader):
img, pert = img.to(config.device), pert.to(config.device)
img_pert = img.clone() + pert
pred = net(img_pert)
pred[pred > 0] = 1.
pred[pred < 1.] = 0.
label[label > 0.5] = 1.
label[label < 1.] = 0.
iou = get_iou(pred[0].data.cpu().numpy(),
label.data.cpu().numpy())
iou_list.append(iou)
iou_list = np.array(iou_list)
print("模型[{}]产生的扰动,模型[{}]的预测结果iou={}", index, index_2,
iou_list.mean())
iou_total.append(iou_list)
iou_total = np.array(iou_total)
np.save("./iou_total.npy", iou_total)
def main():
# 参数
args = get_args()
if not osp.exists(args.result_dir):
os.makedirs(args.result_dir)
print("Evaluating configuration:")
for arg in vars(args):
print("{}:\t{}".format(arg, getattr(args, arg)))
with open('eval-config.json', 'w') as f:
json.dump(args.__dict__, f, indent=4)
# 数据
if args.test:
dataset = SpineDataset(root=args.root,
split='test',
transform=test_transform)
else:
dataset = SpineDataset(root=args.root,
split='val',
transform=val_transform)
dataloader = DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=4,
pin_memory=True)
# 模型
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
if torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
if args.network == 'DeepLab':
model = gcv.models.DeepLabV3(nclass=args.num_classes,
backbone=args.backbone)
elif args.network == 'FCN':
model = gcv.models.FCN(nclass=args.num_classes, backbone=args.backbone)
elif args.network == 'PSPNet':
model = gcv.models.PSP(nclass=args.num_classes, backbone=args.backbone)
elif args.network == 'UNet':
model = UNet(n_class=args.num_classes, backbone=args.backbone)
print('load model from {} ...'.format(args.model))
model.load_state_dict(
torch.load(args.model, map_location='cpu')['state_dict'])
model = model.to(device)
print('Done!')
# 测试
def eval():
with torch.no_grad():
model.eval()
result = []
tq = tqdm.tqdm(total=len(dataloader))
if args.test:
tq.set_description('test')
for i, (data, img_file) in enumerate(dataloader):
tq.update(1)
data = data.to(device)
predict = np.zeros(
(data.size()[1], data.size()[3], data.size()[4]),
dtype=np.uint16)
for idx in range(data.size()[1]):
if args.network in ['DeepLab', 'FCN', 'PSPNet']:
final_out = model(data[:, idx])[0]
elif args.network == 'UNet':
final_out = model(data[:, idx])
predict[idx] = final_out.argmax(
dim=1).cpu().squeeze().numpy().astype(np.uint16)
pred_img = sitk.GetImageFromArray(predict)
test_img = sitk.ReadImage(
osp.join(args.root, 'test', 'image', img_file[0]))
pred_img.CopyInformation(test_img)
result_file = 'mask_' + img_file[0].lower()
sitk.WriteImage(pred_img,
osp.join(args.result_dir, result_file))
else:
tq.set_description('val')
for i, (data, mask, mask_file) in enumerate(dataloader):
tq.update(1)
gt_img = sitk.ReadImage(
osp.join(args.root, 'val', 'groundtruth',
mask_file[0]))
data = data.to(device)
predict = np.zeros(
(data.size()[1], data.size()[3], data.size()[4]),
dtype=np.uint16)
for idx in range(data.size()[1]):
if args.network in ['DeepLab', 'FCN', 'PSPNet']:
final_out = model(data[:, idx])[0]
elif args.network == 'UNet':
final_out = model(data[:, idx])
predict[idx] = final_out.argmax(
dim=1).cpu().squeeze().numpy().astype(np.uint16)
pred_img = sitk.GetImageFromArray(predict)
pred_img.CopyInformation(gt_img)
sitk.WriteImage(pred_img,
osp.join(args.result_dir, mask_file[0]))
ppv, sensitivity, dice, _ = metrics.precision_recall_fscore_support(
mask.numpy().flatten(),
predict.flatten(),
average='binary')
result.append([dice, ppv, sensitivity])
result = np.array(result)
result_mean = result.mean(axis=0)
result_std = result.std(axis=0)
print(result_mean, result_std)
np.savetxt(osp.join(args.result_dir, 'result.txt'),
result_mean,
fmt='%.3f',
header='Dice, Sensitivity, PPV')
tq.close()
eval()
