| 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
|
|
|
