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)
    
    # Determine the result based on the label
    if label == 1:
        result = "AI image"
    else:
        result = "Real image"
        
    return result, f"Confidence: {confidence:.2f}%"

# Define the Gradio interface with updated API
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Textbox(), gr.Textbox()],
    title="Vision Transformer Model",
    description="Upload an image to classify it using the Vision Transformer model.",
    load=gr.Video("load_screen.mp4")  # Specify the loading video
)

# Launch the Gradio interface

demo.launch()