|
|
import numpy as np
|
|
|
import onnxruntime
|
|
|
from huggingface_hub import hf_hub_download
|
|
|
from PIL import Image
|
|
|
|
|
|
|
|
|
def upscale_edsr_2x(image_path: str):
|
|
|
input_image = Image.open(image_path).convert("RGB")
|
|
|
input_image = np.array(input_image).astype("float32")
|
|
|
input_image = np.transpose(input_image, (2, 0, 1))
|
|
|
img_arr = np.expand_dims(input_image, axis=0)
|
|
|
|
|
|
if np.max(img_arr) > 256:
|
|
|
max_range = 65535
|
|
|
else:
|
|
|
max_range = 255.0
|
|
|
img = img_arr / max_range
|
|
|
|
|
|
model_path = hf_hub_download(
|
|
|
repo_id="rupeshs/edsr-onnx",
|
|
|
filename="edsr_onnxsim_2x.onnx",
|
|
|
)
|
|
|
sess = onnxruntime.InferenceSession(model_path)
|
|
|
|
|
|
input_name = sess.get_inputs()[0].name
|
|
|
output_name = sess.get_outputs()[0].name
|
|
|
output = sess.run(
|
|
|
[output_name],
|
|
|
{input_name: img},
|
|
|
)[0]
|
|
|
|
|
|
result = output.squeeze()
|
|
|
result = result.clip(0, 1)
|
|
|
image_array = np.transpose(result, (1, 2, 0))
|
|
|
image_array = np.uint8(image_array * 255)
|
|
|
upscaled_image = Image.fromarray(image_array)
|
|
|
return upscaled_image
|
|
|
|
