import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import os
import argparse
import numpy as np
import cv2
import logging
import math
import _init_paths
from libs.network import U_Net, U_NetDF
from libs.datasets import AcdcDataset
import libs.datasets.joint_augment as joint_augment
import libs.datasets.augment as standard_augment
from libs.datasets.collate_batch import BatchCollator
from libs.configs.config_acdc import cfg
from train_utils.train_utils import load_checkpoint
from utils.metrics import dice
from utils.vis_utils import mask2png, masks_to_contours, apply_mask, img_mask_png
parser = argparse.ArgumentParser(description="arg parser")
parser.add_argument('--used_df', type=str, default=False, help='whether to use df')
parser.add_argument('--selfeat', action='store_true', default=False, help='whether to use feature select')
parser.add_argument('--mgpus', type=str, default=None, required=True, help='whether to use multiple gpu')
parser.add_argument('--model_path1', type=str, default=None, help='whether to train with evaluation')
parser.add_argument('--model_path2', type=str, default=None, help='whether to train with evaluation')
parser.add_argument('--batch_size', type=int, default=1, help='batch size for evaluation')
parser.add_argument('--workers', type=int, default=4, help='number of workers for dataloader')
parser.add_argument('--output_dir', type=str, default=None, required=True, help='specify an output directory if needed')
parser.add_argument('--log_file', type=str, default="../log_evalation.txt", help="the file to write logging")
parser.add_argument('--vis', action='store_true', default=False, help="weather to save test result images")
args = parser.parse_args()
if args.mgpus is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = args.mgpus
def create_logger(log_file):
log_format = '%(asctime)s %(levelname)5s %(message)s'
logging.basicConfig(level=logging.DEBUG, format=log_format, filename=log_file)
console = logging.StreamHandler()
console.setLevel(logging.DEBUG)
console.setFormatter(logging.Formatter(log_format))
logging.getLogger(__name__).addHandler(console)
return logging.getLogger(__name__)
def create_dataloader():
eval_transform = joint_augment.Compose([
joint_augment.To_PIL_Image(),
joint_augment.FixResize(256),
joint_augment.To_Tensor()])
evalImg_transform = standard_augment.Compose([
standard_augment.normalize([cfg.DATASET.MEAN], [cfg.DATASET.STD])])
if cfg.DATASET.NAME == "acdc":
test_set = AcdcDataset(cfg.DATASET.TEST_LIST, df_used=True, joint_augment=eval_transform,
augment=evalImg_transform)
test_loader = DataLoader(test_set, batch_size=1, pin_memory=True,
num_workers=args.workers, shuffle=False,
collate_fn=BatchCollator(size_divisible=32, df_used=True))
return test_loader, test_set
def cal_perfer(preds, masks, tb_dict):
LV_dice = [] # 1
MYO_dice = [] # 2
RV_dice = [] # 3
for i in range(preds.shape[0]):
LV_dice.append(dice(preds[i,1,:,:],masks[i,1,:,:]))
RV_dice.append(dice(preds[i, 3, :, :], masks[i, 3, :, :]))
MYO_dice.append(dice(preds[i, 2, :, :], masks[i, 2, :, :]))
# LV_dice.append(dice(preds[i, 3,:,:],masks[i,1,:,:]))
# RV_dice.append(dice(preds[i, 1, :, :], masks[i, 3, :, :]))
# MYO_dice.append(dice(preds[i, 2, :, :], masks[i, 2, :, :]))
tb_dict["LV_dice"].append(np.mean(LV_dice))
tb_dict["RV_dice"].append(np.mean(RV_dice))
tb_dict["MYO_dice"].append(np.mean(MYO_dice))
return np.mean(LV_dice), np.mean(RV_dice), np.mean(MYO_dice)
def make_one_hot(input, num_classes):
"""Convert class index tensor to one hot encoding tensor.
Args:
input: A tensor of shape [N, 1, *]
num_classes: An int of number of class
Returns:
A tensor of shape [N, num_classes, *]
"""
shape = np.array(input.shape)
shape[1] = num_classes
shape = tuple(shape)
result = torch.zeros(shape).scatter_(1, input.cpu().long(), 1)
# result = result.scatter_(1, input.cpu(), 1)
return result
def test_it(model, data, device="cuda"):
model.eval()
imgs, gts = data[:2]
gts_df = data[2]
imgs = imgs.to(device)
gts = gts.to(device)
net_out = model(imgs)
if len(net_out) > 1:
preds_out = net_out[0]
preds_df = net_out[1]
else:
preds_out = net_out[0]
preds_df = None
preds_out = nn.functional.softmax(preds_out, dim=1)
_, preds = torch.max(preds_out, 1)
preds = preds.unsqueeze(1) # (N, 1, *)
return preds, preds_df
def vis_it(pred, gt, img=None, filename=None, infos=None):
h, w = pred.shape
# gt_contours = masks_to_contours(gt)
# mask = np.hstack([pred, np.zeros((h, 1)), gt])
# gt_contours = np.hstack([gt_contours, np.zeros((h, 1)), np.zeros_like(gt)])
# im_rgb = mask2png(mask).astype(np.int16)
# im_rgb[:, w, :] = [255, 255, 255]
# im_rgb = apply_mask(im_rgb, gt_contours, [255, 255, 255], 0.8)
pred_im = mask2png(pred).astype(np.int16)
gt_im = mask2png(gt).astype(np.int16)
img = (img - img.min()) / (img.max() - img.min()) * 255
img = np.stack([img, img, img], axis=2)
# img = img_mask_png(img, gt, alpha=0.1)
# im_rgb = np.hstack([im_rgb, 255*np.ones((h, 1, 3)), img])
# cv2.putText(im_rgb, "prediction", (2,h-4),
# fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.4, color=(255, 255, 255), thickness=1)
# cv2.putText(im_rgb, "ground truth", (w, h-4),
# fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.4, color=(255, 255, 255), thickness=1)
# st_pos = 15
# if infos is not None:
# for info in infos:
# cv2.putText(im_rgb, info+": {}".format(infos[info]), (2, st_pos),
# fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.4, color=(255, 255, 255), thickness=1)
# st_pos += 10
cv2.imwrite(filename+"_img.png", img[:,:,::-1])
cv2.imwrite(filename+"_pred.png", pred_im[:,:,::-1])
cv2.imwrite(filename+"_gt.png", gt_im[:,:,::-1])
def vis_df(pred_df, gt_df, filename, infos=None):
_, h, w = pred_df.shape
# save .npy files
np.save(filename+'.npy', [pred_df, gt_df])
theta = np.arctan2(gt_df[1,...], gt_df[0,...])
degree_gt = (theta - theta.min()) / (theta.max() - theta.min()) * 255
# degree_gt = theta * 360
mag_gt = np.sum(gt_df ** 2, axis=0, keepdims=False)
mag_gt = mag_gt / mag_gt.max() * 255
theta = np.arctan2(pred_df[1,...], pred_df[0,...])
degree_df = (theta - theta.min()) / (theta.max() - theta.min()) * 255
# degree_df = theta * 360
magnitude = np.sum(pred_df ** 2, axis=0, keepdims=False)
magnitude = magnitude / magnitude.max() * 255
im = np.hstack([magnitude, np.zeros((h, 1)), mag_gt, np.zeros((h, 1)), degree_df, np.zeros((h, 1)), degree_gt]).astype(np.uint8)
im = cv2.applyColorMap(im, cv2.COLORMAP_JET)
cv2.imwrite(filename+"_df_pred_mag.png", im[:h, :w, ...])
cv2.imwrite(filename+"_df_gt_mag.png", im[:h, w+1:2*w+1, ...])
cv2.imwrite(filename+"_df_pred_degree.png", im[:h, 2*w+2:3*w+2, ...])
cv2.imwrite(filename+"_df_gt_degree.png", im[:h, 3*w+3:, ...])
def test():
root_result_dir = args.output_dir
os.makedirs(root_result_dir, exist_ok=True)
log_file = os.path.join(root_result_dir, args.log_file)
logger = create_logger(log_file)
for key, val in vars(args).items():
logger.info("{:16} {}".format(key, val))
# create dataset & dataloader & network
if args.used_df == 'U_NetDF':
model = U_NetDF(selfeat=args.selfeat, num_class=4, auxseg=True)
elif args.used_df == 'U_Net':
model = U_Net(num_class=4)
if args.mgpus is not None and len(args.mgpus) > 2:
model = nn.DataParallel(model)
model.cuda()
test_loader, test_set = create_dataloader()
checkpoint = torch.load(args.model_path1, map_location='cpu')
model.load_state_dict(checkpoint['model_state'])
dice_dict = {"LV_dice": [],
"RV_dice": [],
"MYO_dice": []}
for i, data in enumerate(test_loader):
if i != 23: continue
# i = 5405
# data = test_set[5405]
# data = [data[0][None], data[1][None], data[2][None]]
pred, pred_df = test_it(model, data[:3])
_, gt, gt_df = data[:3]
gt = gt.to("cuda")
L, R, MYO = cal_perfer(make_one_hot(pred, 4), make_one_hot(gt, 4), dice_dict)
data_info = test_set.data_infos[i]
if args.vis:
# if 0.7 <= (L + R + MYO) / 3 < 0.8:
vis_it(pred.cpu().numpy()[0, 0], gt.cpu().numpy()[0, 0], data[0].cpu().numpy()[0, 0],
filename=os.path.join(root_result_dir, str(i)))
if pred_df is not None:
vis_df(pred_df.detach().cpu().numpy()[0], gt_df.cpu().numpy()[0],
filename=os.path.join(root_result_dir, str(i)))
print("\r{}/{} {:.0%} {}".format(i, len(test_set), i/len(test_set),
np.mean(list(dice_dict.values()))), end="")
print()
logger.info("2D Dice Metirc:")
logger.info("Total {}".format(len(test_set)))
logger.info("LV_dice: {}".format(np.mean(dice_dict["LV_dice"])))
logger.info("RV_dice: {}".format(np.mean(dice_dict["RV_dice"])))
logger.info("MYO_dice: {}".format(np.mean(dice_dict["MYO_dice"])))
logger.info("Mean_dice: {}".format(np.mean(list(dice_dict.values()))))
if __name__ == "__main__":
test()
