import torch import pandas as pd import numpy as np from common.utils import PSNR, cal_ssim, logger_info, _rgb2ycbcr, modcrop from pathlib import Path from PIL import Image import time def val_image_pair(model, hr_image, lr_image, output_image_path=None, device='cuda'): with torch.inference_mode(): start_time = time.perf_counter_ns() # prepare lr_image lr_image = torch.tensor(lr_image).type(torch.float32).permute(2,0,1) lr_image = lr_image.unsqueeze(0).to(torch.device(device)) # predict pred_lr_image = model(lr_image) # postprocess pred_lr_image = pred_lr_image.squeeze(0).permute(1,2,0).type(torch.uint8) pred_lr_image = pred_lr_image.cpu().numpy() run_time_ns = time.perf_counter_ns() - start_time torch.cuda.empty_cache() if not output_image_path is None: Image.fromarray(pred_lr_image).save(output_image_path) # metrics hr_image = modcrop(hr_image, model.scale) Y_left, Y_right = _rgb2ycbcr(pred_lr_image)[:, :, 0], _rgb2ycbcr(hr_image)[:, :, 0] lr_area = np.prod(lr_image.shape[-2:]) return PSNR(Y_left, Y_right, model.scale), cal_ssim(Y_left, Y_right), run_time_ns, lr_area def valid_steps(model, datasets, config, log_prefix=""): dataset_names = list(datasets.keys()) results = [] for i in range(len(dataset_names)): 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(): predictions_path.mkdir() test_dataset = datasets[dataset_name] tasks = [] for hr_image, lr_image, hr_image_path, lr_image_path in test_dataset: output_image_path = predictions_path / f'{Path(hr_image_path).stem}.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 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 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, 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("\n" + str(pd.DataFrame([row], columns=column_names).set_index('Dataset').T)) config.writer.flush() results = pd.DataFrame(results, columns=column_names).set_index('Dataset') return results