app.py

# -*- coding: utf-8 -*-
"""HED.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1sxAQYKOi_hJozIVNX80ChnznWExGtWXd

# IMPORTS
"""
import tensorflow as tf
import onnx
import gradio as gr

import onnx
model = onnx.load("vit_onnx.onnx")

import onnxruntime as ort
session = ort.InferenceSession("vit_onnx.onnx")

input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

CLASS_NAMES = ["Angry", "Happy", "Sad"]
def predict_image(im):
  im = tf.expand_dims(tf.cast(im, tf.float32), axis=0).numpy()
  prediction = session.run([output_name], {input_name: im})
  return {CLASS_NAMES[i]: float(prediction[0][0][i]) for i in range(3)}

image = gr.inputs.Image(shape=(224, 224))
label = gr.outputs.Label(num_top_classes=3)
iface = gr.Interface(fn=predict_image, inputs=image, outputs=label, capture_session=True)
iface.launch()