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