import gradio as gr

# Define the function to predict using the model
def predict_image(img):
    model = gr.load("models/google/vit-base-patch16-224")
    return model.predict(img)

# Create the interface
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(),
    outputs="label",
    live=True,
    capture_session=True,
    title="Image recognition",
    description="Upload an image you want to categorize.",
    theme="Monochrome"
)

iface.launch()