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