Vladimir Protsenko 8 months ago
parent ba19c2eee0
commit e5a4c9531a

@ -29,7 +29,8 @@ def val_image_pair(model, hr_image, lr_image, output_image_path=None, device='cu
# metrics
hr_image = modcrop(hr_image, model.scale)
left, right = _rgb2ycbcr(pred_lr_image)[:, :, 0], _rgb2ycbcr(hr_image)[:, :, 0]
return PSNR(left, right, model.scale), cal_ssim(left, right), run_time_ns
lr_area = np.prod(lr_image.shape[-2:])
return PSNR(left, right, model.scale), cal_ssim(left, right), run_time_ns, lr_area
def valid_steps(model, datasets, config, log_prefix=""):
# ray.init(num_cpus=16, num_gpus=1, ignore_reinit_error=True, log_to_driver=False, runtime_env={"working_dir": "../"})
@ -40,6 +41,9 @@ def valid_steps(model, datasets, config, log_prefix=""):
dataset_name = dataset_names[i]
psnrs, ssims = [], []
run_times_ns = []
lr_areas = []
total_area = 0
start_time = time.time()
predictions_path = config.valout_dir / dataset_name
if not predictions_path.exists():
@ -51,6 +55,8 @@ def valid_steps(model, datasets, config, log_prefix=""):
output_image_path = predictions_path / f'{Path(hr_image_path).stem}_rcnet.png' if config.save_predictions else None
task = val_image_pair(model, hr_image, lr_image, output_image_path, device=config.device)
tasks.append(task)
total_time = time.time() - start_time
# ready_refs, remaining_refs = ray.wait(tasks, num_returns=1, timeout=None)
# while len(remaining_refs) > 0:
# print(f"\rReady {len(ready_refs)+1}/{len(test_dataset)}", end=" ")
@ -58,22 +64,42 @@ def valid_steps(model, datasets, config, log_prefix=""):
# print("\r", end=" ")
# tasks = [ray.get(task) for task in tasks]
for psnr, ssim, run_time_ns in tasks:
for psnr, ssim, run_time_ns, lr_area in tasks:
psnrs.append(psnr)
ssims.append(ssim)
run_times_ns.append(run_time_ns)
lr_areas.append(lr_area)
total_area += lr_area
config.logger.info(
'\r{} | Dataset {} | AVG Val PSNR: {:02f}, AVG: SSIM: {:04f}'.format(log_prefix, dataset_name, np.mean(np.asarray(psnrs)), np.mean(np.asarray(ssims))))
config.writer.flush()
results.append([
row = [
dataset_name,
np.mean(psnrs),
np.mean(ssims),
np.mean(run_times_ns)*1e-9,
np.percentile(run_times_ns, q=95)*1e-9])
np.percentile(run_times_ns, q=95)*1e-9,
len(test_dataset),
np.mean(lr_areas),
total_area,
total_time
]
results.append(row)
column_names = [
'Dataset',
'AVG PSNR',
'AVG SSIM',
f'AVG {config.device} Time, s',
f'P95 {config.device} Time, s',
'Image count',
'AVG image area',
'Total area',
'Total time, s'
]
config.logger.info(pd.DataFrame(
[row],
columns=column_names
).set_index('Dataset'))
config.writer.flush()
results = pd.DataFrame(results, columns=['Dataset', 'PSNR', 'SSIM', f'AVG {config.device} Time, s', f'P95 {config.device} Time, s']).set_index('Dataset')
results = pd.DataFrame(results, columns=column_names).set_index('Dataset')
return results

@ -40,7 +40,7 @@ class ValOptions():
if not args.valout_dir.exists():
args.valout_dir.mkdir()
args.current_iter = args.model_name.split('_')[-1]
args.results_path = os.path.join(args.valout_dir, f'results_{args.device}.csv')
args.results_path = os.path.join(args.valout_dir, f'results_{args.model_name}_{args.device}.csv')
# Tensorboard for monitoring
writer = SummaryWriter(log_dir=args.valout_dir)
logger_name = f'val_{args.model_path.stem}'

Loading…
Cancel
Save