|
import torch |
|
import os |
|
import sys |
|
import pathlib |
|
CURRENT_DIR = pathlib.Path(__file__).parent |
|
sys.path.append(str(CURRENT_DIR)) |
|
from tqdm import tqdm |
|
import data |
|
import metric |
|
import onnxruntime |
|
import cv2 |
|
from data.data_tiling import tiling_inference |
|
import argparse |
|
|
|
class Configs(): |
|
def __init__(self): |
|
parser = argparse.ArgumentParser(description='SR') |
|
|
|
|
|
parser.add_argument('--ipu', action='store_true', |
|
help='use ipu') |
|
parser.add_argument('--onnx_path', type=str, default='RCAN_int8_NHWC.onnx', |
|
help='onnx path') |
|
parser.add_argument('--provider_config', type=str, default=None, |
|
help='provider config path') |
|
|
|
parser.add_argument('--dir_data', type=str, default='dataset/', |
|
help='dataset directory') |
|
parser.add_argument('--data_test', type=str, default='Set5', |
|
help='test dataset name') |
|
|
|
parser.add_argument('--n_threads', type=int, default=6, |
|
help='number of threads for data loading') |
|
parser.add_argument('--scale', type=str, default='2', |
|
help='super resolution scale, now only support x2') |
|
self.parser = parser |
|
|
|
def parse(self): |
|
args = self.parser.parse_args() |
|
args.scale = list(map(lambda x: int(x), args.scale.split('+'))) |
|
args.data_test = args.data_test.split('+') |
|
print(args) |
|
return args |
|
|
|
|
|
|
|
def quantize(img, rgb_range): |
|
pixel_range = 255 / rgb_range |
|
return img.mul(pixel_range).clamp(0, 255).round().div(pixel_range) |
|
|
|
def test_model(session, loader, device): |
|
torch.set_grad_enabled(False) |
|
self_scale = [2] |
|
for idx_data, d in enumerate(loader.loader_test): |
|
eval_ssim = 0 |
|
eval_psnr = 0 |
|
for idx_scale, scale in enumerate(self_scale): |
|
d.dataset.set_scale(idx_scale) |
|
for lr, hr, filename in tqdm(d, ncols=80): |
|
sr = tiling_inference(session, lr.cpu().numpy(), 8, (56, 56)) |
|
sr = torch.from_numpy(sr).to(device) |
|
sr = quantize(sr, 255) |
|
eval_psnr += metric.calc_psnr( |
|
sr, hr, scale, 255, benchmark=d) |
|
eval_ssim += metric.calc_ssim( |
|
sr, hr, scale, 255, dataset=d) |
|
mean_ssim = eval_ssim / len(d) |
|
mean_psnr = eval_psnr / len(d) |
|
print("psnr: %s, ssim: %s"%(mean_psnr, mean_ssim)) |
|
return mean_psnr, mean_ssim |
|
|
|
def main(args): |
|
loader = data.Data(args) |
|
if args.ipu: |
|
providers = ["VitisAIExecutionProvider"] |
|
provider_options = [{"config_file": args.provider_config}] |
|
else: |
|
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] |
|
provider_options = None |
|
onnx_file_name = args.onnx_path |
|
ort_session = onnxruntime.InferenceSession(onnx_file_name, providers=providers, provider_options=provider_options) |
|
test_model(ort_session, loader, device="cpu") |
|
|
|
|
|
if __name__ == '__main__': |
|
args = Configs().parse() |
|
main(args) |
|
|
