import os import copy import time import numpy as np import onnxruntime from PIL import Image, ImageOps import gradio def run_inference(onnx_session, input_size, image): # Resize temp_image = image.copy() resize_image = temp_image.resize((input_size, input_size), Image.ANTIALIAS) x = ImageOps.exif_transpose(resize_image) x = np.array(x) # Preprocessing x = x.astype(np.float32) mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) std = np.array([0.229, 0.224, 0.225], dtype=np.float32) x = (x / 255 - mean) / std x = x.transpose(2, 0, 1).astype('float32') x = x.reshape(-1, 3, input_size, input_size) # Inference input_name = onnx_session.get_inputs()[0].name output_name = onnx_session.get_outputs()[0].name onnx_result = onnx_session.run([output_name], {input_name: x}) # Postprocessing onnx_result = np.array(onnx_result).squeeze() min_value = np.min(onnx_result) max_value = np.max(onnx_result) onnx_result = (onnx_result - min_value) / (max_value - min_value) onnx_result *= 255 onnx_result = onnx_result.astype('uint8') return onnx_result # Load model onnx_session = onnxruntime.InferenceSession("u2net.onnx") def create_rgba(mode, image): image = Image.fromarray(image).convert('RGB') out = run_inference( onnx_session, 320, image, ) resize_image = Image.fromarray(out).resize((image.size[0], image.size[1]), Image.ANTIALIAS) if mode == "binary": resize_image = resize_image.point(lambda x: 255 if x > 125 else 0) mask = resize_image rgba_image = image.convert('RGBA') rgba_image.putalpha(mask) return rgba_image inputs = [gradio.inputs.Radio(["binary", "smooth"]), gradio.inputs.Image()] outputs = gradio.outputs.Image(type='numpy') gradio.Interface(fn=create_rgba, inputs=inputs, outputs=outputs).launch()