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

# Initialize the model
model = CustomModel()

def predict(image: Image.Image):
        animation.visible = True  # Show the animation
        label, confidence = model.predict(image)
        result = "AI image" if label == 1 else "Real image"
        animation.visible = False  # Hide the animation
        return result, f"Confidence: {confidence:.2f}%"

theme = gr.themes.Base()

with gr.Blocks(theme=theme) as demo:
    # Define the Gradio interface
    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.",
        live=True,  # Allows immediate prediction
    )

# Launch the Gradio interface
if __name__ == "__main__":
    demo.launch()