class ADNTester(Tester):
def __init__(self, **params):
super(ADNTester, self).__init__(**params)
def get_image(self, data):
data_opts = self.opts.dataset
dataset_type = data_opts["dataset_type"]
if dataset_type == "deep_lesion":
if data_opts[dataset_type]["load_mask"]:
return data['lq_image'], data['hq_image'], data[
"data_name"], data["mask"]
else:
return data['lq_image'], data['hq_image'], data["data_name"]
elif dataset_type == "spineweb":
return data['a'], data['b'], data["data_name"]
elif dataset_type == "nature_image":
return data['artifact'], data['no_artifact'], data["data_name"]
def get_metric(self, metric):
def measure(x, y):
x = self.dataset.to_numpy(x, False)
y = self.dataset.to_numpy(y, False)
x = x * 0.5 + 0.5
y = y * 0.5 + 0.5
return metric(x, y, data_range=1.0)
return measure
def get_pairs(self):
if hasattr(self.model, 'mask'):
mask = self.model.mask
img_low = self.model.img_low * mask
img_high = self.model.img_high * mask
pred_lh = self.model.pred_lh * mask
else:
img_low = self.model.img_low
img_high = self.model.img_high
pred_lh = self.model.pred_lh
return [("before", (img_low, img_high)),
("after", (pred_lh, img_high))], self.model.name
def get_visuals(self, n=8):
lookup = [("l", "img_low"), ("ll", "pred_ll"), ("lh", "pred_lh"),
("h", "img_high"), ("hl", "pred_hl"), ("hh", "pred_hh")]
visual_window = self.opts.visual_window
def visual_func(x):
x = x * 0.5 + 0.5
x[x < visual_window[0]] = visual_window[0]
x[x > visual_window[1]] = visual_window[1]
x = (x - visual_window[0]) / (visual_window[1] - visual_window[0])
return x
return self.model._get_visuals(lookup, n, visual_func, False)
def get_logger(self, opts):
self.logger = Logger(self.run_dir, self.epoch, self.run_name)
self.logger.add_iter_visual_log(self.get_visuals, 1, "test_visuals")
self.logger.add_metric_log(self.get_pairs,
(("ssim", self.get_metric(ssim)),
("psnr", self.get_metric(psnr))),
opts.metrics_step)
return self.logger
def evaluate(self, model, data):
model.evaluate(*data[:3])
if len(data) == 4:
mask = 1 - data[3].to(model.img_high)
model.mask = mask
dataset_opts = opts['dataset']
train_dataset = get_dataset(**dataset_opts)
train_loader = DataLoader(train_dataset,
batch_size=opts["batch_size"], num_workers=opts['num_workers'], shuffle=True)
train_loader = add_post(train_loader, get_image)
# Get checkpoint
if opts['last_epoch'] == 'last':
checkpoint, start_epoch = get_last_checkpoint(run_dir)
else:
start_epoch = opts['last_epoch']
checkpoint = path.join(run_dir, "net_{}".format(start_epoch))
if type(start_epoch) is not int: start_epoch = 0
# Get model
model = ADNTrain(opts['learn'], opts['loss'], **opts['model'])
if opts['use_gpu']: model.cuda()
if path.isfile(checkpoint): model.resume(checkpoint)
# Get logger
logger = Logger(run_dir, start_epoch, args.run_name)
logger.add_loss_log(model.get_loss, opts["print_step"], opts['window_size'])
logger.add_iter_visual_log(model.get_visuals, opts['visualize_step'], "train_visuals")
logger.add_save_log(model.save, opts['save_step'])
# Train the model
for epoch in range(start_epoch, opts['num_epochs']):
for data in logger(train_loader):
model.optimize(*data)
model.update_lr()