Spaces:
Running
Running
File size: 1,901 Bytes
6d8a73e e28e2c8 6d8a73e e28e2c8 6d8a73e e28e2c8 6d8a73e e28e2c8 6d8a73e e28e2c8 6d8a73e e28e2c8 6d8a73e e28e2c8 6d8a73e e28e2c8 6d8a73e e28e2c8 6d8a73e e28e2c8 6d8a73e 80a9211 6d8a73e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
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()
|