0xrk

Can someone help me providing the modified code for exporting scale arg

​

import io

import numpy as np

import onnx

import onnxruntime as ort

import shutil

​

from PIL import Image

from pathlib import Path

from onnxruntime_extensions.tools.pre_post_processing import *

​

# NOTE: Currently assumes we're running from the root directory of the Real-ESRGAN repo so the weights and input image

# are loaded using paths relative to that.

_this_dirpath = Path(__file__).parent

​

def create_onnx_model(onnx_model_path: str):

import torch

from RealESRGAN import RealESRGAN

​

weights_path = 'weights/RealESRGAN_x4.pth'

if not Path(weights_path).exists():

raise ValueError(f"{weights_path} not found. Please download the RealESRGAN_x4.pth weights as per "

"https://github.com/ai-forever/Real-ESRGAN/blob/main/weights/README.md")

​

device = torch.device('cpu')

esrgan_model = RealESRGAN(device, scale=4)

esrgan_model.load_weights(weights_path, download=True)

​

# We export the torch model in the RealESRGAN \model` property`

torch_model = esrgan_model.model

​

# set the model to inference mode

torch_model.eval()

​

# Create random input to the model

x = torch.randn(1, 3, 240, 240)

​

# Export the model

torch.onnx.export(torch_model, # model being run

(x, ), # model input (or a tuple for multiple inputs)

onnx_model_path, # where to save the model (can be a file or file-like object)

export_params=True, # store the trained parameter weights inside the model file

opset_version=15, # the ONNX version to export the model to

do_constant_folding=True, # whether to execute constant folding for optimization

input_names=['input'], # the model's input names

output_names=['output']) # the model's output names

​

​

def add_pre_post_processing(input_model_path: str, output_model_path: str):

# we do a Resize with anti-aliasing which requires ONNX opset 18 in onnxruntime version 1.14 or later

from packaging import version

if version.parse(ort.__version__) < version.parse("1.14.0"):

raise ValueError("ONNX Runtime version 1.14 or later required. Please update your onnxruntime python package.")

​

onnx_opset = 18

model = onnx.load(input_model_path)

inputs = [create_named_value("image", onnx.TensorProto.UINT8, ["num_bytes"])]

​

# infer the input sizes from the model.

model_input_shape = model.graph.input[0].type.tensor_type.shape

w_in = model_input_shape.dim[-1].dim_value

h_in = model_input_shape.dim[-2].dim_value

​

pipeline = PrePostProcessor(inputs, onnx_opset)

pipeline.add_pre_processing(

[

ConvertImageToBGR(), # jpg/png image to BGR in HWC layout

ReverseAxis(axis=2, dim_value=3, name="BGR_to_RGB"), # BGR to RGB

Resize((h_in, w_in), layout='HWC'),

CenterCrop(h_in, w_in), # CenterCrop requires HWC

ChannelsLastToChannelsFirst(), # HWC to CHW

ImageBytesToFloat(), # convert to float in range 0..1

Unsqueeze([0]), # add batch dimensions

]

)

​

pipeline.add_post_processing(

[

Squeeze([0]), # remove batch dimensions

FloatToImageBytes(), # convert back to uint8

Transpose(perms=[1, 2, 0], name="CHW_to_HWC"), # channels first to channels last

ReverseAxis(axis=2, dim_value=3, name="RGB_to_BGR"), # RGB to BGR

ConvertBGRToImage(image_format="png", name="convert_to_png"),

]

)

​

new_model = pipeline.run(model)

onnx.save_model(new_model, output_model_path)

​

​

def test_onnx_model(model_path: str):

from onnxruntime_extensions import get_library_path

​

so = ort.SessionOptions()

so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED

​

# register the custom operators for the image decode/encode pre/post processing provided by onnxruntime-extensions

# with onnxruntime. if we do not do this we'll get an error on model load about the operators not being found.

ortext_lib_path = get_library_path()

so.register_custom_ops_library(ortext_lib_path)

inference_session = ort.InferenceSession(model_path, so)

​

test_image_path = _this_dirpath / 'MauiSuperResolution' / 'Resources' / 'Raw' / 'lr_lion.png'

test_image_bytes = np.fromfile(test_image_path, dtype=np.uint8)

outputs = inference_session.run(['image_out'], {'image': test_image_bytes})

upsized_image_bytes = outputs[0]

​

original_img = Image.open(io.BytesIO(test_image_bytes))

updated_img = Image.open(io.BytesIO(upsized_image_bytes))

​

# centered crop of original to match the area processed

def _center_crop_to_square(img: Image):

if img.height != img.width:

target_size = img.width if img.width < img.height else img.height

w_start = int(np.floor((img.width - target_size) / 2))

w_end = w_start + target_size

h_start = int(np.floor((img.height - target_size) / 2))

h_end = h_start + target_size

​

return img.crop((w_start, h_start, w_end, h_end))

else:

return img

​

original_cropped_img = _center_crop_to_square(original_img)

​

new_width, new_height = updated_img.size

​

# create a side-by-side image with both.

# resize the original to the model input size followed by the output size so the processing is more equivalent

original_cropped_img = original_cropped_img.resize((240, 240))

resized_orig_img = original_cropped_img.resize((new_width, new_height))

​

combined = Image.new('RGB', (new_width * 2, new_height))

combined.paste(resized_orig_img, (0, 0))

combined.paste(updated_img, (new_width, 0))

combined.show('Original resized vs Super Resolution resized')

# combined.save('Original resized vs Super Resolution resized.png', format='PNG')

​

​

def main():

onnx_model_path = Path('RealESRGAN.onnx')

if not onnx_model_path.exists():

print("Creating ONNX model from pytorch model...")

create_onnx_model(str(onnx_model_path))

​

assert onnx_model_path.exists()

​

print("Adding pre/post processing to ONNX model...")

output_model_path = str(onnx_model_path).replace(".onnx", "_with_pre_post_processing.onnx")

add_pre_post_processing(str(onnx_model_path), output_model_path)

​

print("Testing ONNX model with pre/post processing")

test_onnx_model(output_model_path)

​

print("Copying ONNX model to MAUI applications Resources/Raw directory...")

shutil.copy("RealESRGAN_with_pre_post_processing.onnx",

_this_dirpath / 'MauiSuperResolution' / 'Resources' / 'Raw')

​

if __name__ == '__main__':

main()