import gradio as gr
import numpy as np
import tensorflow as tf
from huggingface_hub import from_pretrained_keras

description = "Keras implementation for Video Vision Transformer to classify samples of medmnist"
article = "Author:<a href=\"https://huggingface.co/pablorodriper\"> Pablo Rodríguez</a>; Based on the keras example by <a href=\"https://keras.io/examples/vision/vivit/\">Aritra Roy Gosthipaty and Ayush Thakur</a>"
title = "Video Vision Transformer on medmnist"

def infer(x):
    return model.predict(tf.expand_dims(x, axis=0))[0]

model = from_pretrained_keras("pablorodriper/vivit")

iface = gr.Interface(
    fn = infer,
    inputs = "video",
    outputs = "number",
    description = description,
    title = title,
    article = article
    )

iface.launch()