File size: 1,361 Bytes
35752ce b0ea2eb 35752ce b0ea2eb |
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 |
import numpy as np
import onnxruntime as ort
import gradio as gr
from PIL import Image
from torchvision.models import ResNet50_Weights
weights = ResNet50_Weights.DEFAULT
preprocess = weights.transforms() # Necessary input transformations
ort_session = ort.InferenceSession("resnet50.onnx", providers=["CPUExecutionProvider"])
def preprocess_inputs(img: Image):
img = preprocess(img)
img_array = np.array(img).astype(np.float32)
img_array = np.expand_dims(img_array, axis=0)
return img_array
def predict(img):
img = preprocess_inputs(img)
ort_inputs = {ort_session.get_inputs()[0].name: img}
ort_outputs = ort_session.run(None, ort_inputs)
#probs = np.exp(ort_outputs) / np.sum(np.exp(ort_outputs)) # softmax
label_index = np.argmax(ort_outputs[0], axis=1).item()
predicted_label = weights.meta["categories"][label_index]
return predicted_label
demo = gr.Interface(predict, gr.Image(type="pil", image_mode="RGB"), gr.Label(),
title="ResNet-50 Using onnxruntime",
description="Upload any image and see if resnet-50 can classify it! (1000 possible image classes)",
article="Part of a tutorial on [how to deploy an ONNX mode to Hugging Face](https://liamgroen.nl/posts/day-6-deploying-model-to-huggingface-spaces-through-onnx/index.html)")
demo.launch() |