import gradio as gr
from PIL import Image
from vit_model_test import CustomModel

# Initialize the model
model = CustomModel()

def predict(image: Image.Image):
    # Get predictions from the model
    label, confidence = model.predict(image)
    return f"Predicted label: {label}", f"Confidence: {confidence:.2f}%"

# Define the Gradio interface with updated API
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),  # Updated for new Gradio API
    outputs=[gr.Textbox(), gr.Textbox()],  # Updated for new Gradio API
    title="Vision Transformer Model",    # Title of the Gradio interface
    description="Upload an image to classify it using the Vision Transformer model."  # Description
)

# Launch the Gradio interface

demo.launch()