Vladimir Protsenko 6 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 # metrics
hr_image = modcrop(hr_image, model.scale) hr_image = modcrop(hr_image, model.scale)
left, right = _rgb2ycbcr(pred_lr_image)[:, :, 0], _rgb2ycbcr(hr_image)[:, :, 0] 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=""): 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": "../"}) # ray.init(num_cpus=16, num_gpus=1, ignore_reinit_error=True, log_to_driver=False, runtime_env={"working_dir": "../"})
@ -39,7 +40,10 @@ def valid_steps(model, datasets, config, log_prefix=""):
for i in range(len(dataset_names)): for i in range(len(dataset_names)):
dataset_name = dataset_names[i] dataset_name = dataset_names[i]
psnrs, ssims = [], [] psnrs, ssims = [], []
run_times_ns = [] run_times_ns = []
lr_areas = []
total_area = 0
start_time = time.time()
predictions_path = config.valout_dir / dataset_name predictions_path = config.valout_dir / dataset_name
if not predictions_path.exists(): 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 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) task = val_image_pair(model, hr_image, lr_image, output_image_path, device=config.device)
tasks.append(task) tasks.append(task)
total_time = time.time() - start_time
# ready_refs, remaining_refs = ray.wait(tasks, num_returns=1, timeout=None) # ready_refs, remaining_refs = ray.wait(tasks, num_returns=1, timeout=None)
# while len(remaining_refs) > 0: # while len(remaining_refs) > 0:
# print(f"\rReady {len(ready_refs)+1}/{len(test_dataset)}", end=" ") # 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=" ") # print("\r", end=" ")
# tasks = [ray.get(task) for task in tasks] # 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) psnrs.append(psnr)
ssims.append(ssim) ssims.append(ssim)
run_times_ns.append(run_time_ns) run_times_ns.append(run_time_ns)
lr_areas.append(lr_area)
total_area += lr_area
config.logger.info( row = [
'\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([
dataset_name, dataset_name,
np.mean(psnrs), np.mean(psnrs),
np.mean(ssims), np.mean(ssims),
np.mean(run_times_ns)*1e-9, 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 return results

@ -40,7 +40,7 @@ class ValOptions():
if not args.valout_dir.exists(): if not args.valout_dir.exists():
args.valout_dir.mkdir() args.valout_dir.mkdir()
args.current_iter = args.model_name.split('_')[-1] 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 # Tensorboard for monitoring
writer = SummaryWriter(log_dir=args.valout_dir) writer = SummaryWriter(log_dir=args.valout_dir)
logger_name = f'val_{args.model_path.stem}' logger_name = f'val_{args.model_path.stem}'

Loading…
Cancel
Save