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