53 lines
1.4 KiB
Python
53 lines
1.4 KiB
Python
|
from os import path
|
||
|
|
|
||
|
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
import onnxruntime as ort
|
||
|
|
import torch
|
||
|
|
import time
|
||
|
|
|
||
|
|
cfg = 8
|
||
|
|
steps = 22
|
||
|
|
height = 512
|
||
|
|
width = 512
|
||
|
|
|
||
|
|
esrgan = path.join('..', 'models', 'RealESRGAN_x4plus.onnx')
|
||
|
|
output = path.join('..', 'outputs', 'test.png')
|
||
|
|
upscale = path.join('..', 'outputs', 'test-large.png')
|
||
|
|
|
||
|
|
print('upscaling test image...')
|
||
|
|
session = ort.InferenceSession(esrgan, providers=['DmlExecutionProvider'])
|
||
|
|
|
||
|
|
in_image = cv2.imread(output, cv2.IMREAD_UNCHANGED)
|
||
|
|
|
||
|
|
in_mat = cv2.cvtColor(in_image, cv2.COLOR_BGR2RGB)
|
||
|
|
print('shape before', np.shape(in_mat))
|
||
|
|
in_mat = np.transpose(in_mat, (2, 1, 0))[np.newaxis]
|
||
|
|
print('shape after', np.shape(in_mat))
|
||
|
|
in_mat = in_mat.astype(np.float32)
|
||
|
|
in_mat = in_mat/255
|
||
|
|
|
||
|
|
start_time = time.time()
|
||
|
|
input_name = session.get_inputs()[0].name
|
||
|
|
output_name = session.get_outputs()[0].name
|
||
|
|
in_mat = torch.tensor(in_mat)
|
||
|
|
out_mat = session.run([output_name], {
|
||
|
|
input_name: in_mat.cpu().numpy()
|
||
|
|
})[0]
|
||
|
|
elapsed_time = time.time() - start_time
|
||
|
|
print(elapsed_time)
|
||
|
|
|
||
|
|
print('output shape', np.shape(out_mat))
|
||
|
|
out_mat = np.squeeze(out_mat, (0))
|
||
|
|
print(np.shape(out_mat))
|
||
|
|
out_mat = np.transpose(out_mat, (2, 1, 0))
|
||
|
|
print(out_mat, np.shape(out_mat))
|
||
|
|
out_mat = np.clip(out_mat, 0.0, 1.0)
|
||
|
|
out_mat = out_mat * 255
|
||
|
|
out_mat = out_mat.astype(np.uint8)
|
||
|
|
out_image = cv2.cvtColor(out_mat, cv2.COLOR_RGB2BGR)
|
||
|
|
|
||
|
|
cv2.imwrite(upscale, out_image)
|
||
|
|
|
||
|
|
print('upscaled test image to %s' % upscale)
|