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

# Initialize the model
model = CustomModel()

def predict(image: Image.Image):
    try:
        label, confidence = model.predict(image)
        result = "AI image" if label == 1 else "Real image"
        return result, f"Confidence: {confidence:.2f}%"
    except Exception as e:
        return "Error in prediction", str(e)

# Custom HTML and CSS to replace the logo with a video
custom_html = """
<style>
    /* Hide the default logo */
    .gradio-logo {
        display: none;
    }
    /* Center the video container */
    .loading-container {
        text-align: center;
        margin-top: 20px;
    }
    .loading-container video {
        width: 320px; /* Adjust video size */
        height: auto; /* Maintain aspect ratio */
    }
</style>
<div class="loading-container">
    <video autoplay muted loop>
        <source src="load_screen.mp4" type="video/mp4">
        Your browser does not support the video tag.
    </video>
    <div>Processing, please wait...</div>
</div>
"""

# Define the Gradio interface
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."
)

# Inject the custom HTML to show the video instead of the logo
demo.load(custom_html)

# Launch the Gradio interface
demo.launch()