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

# Initialize the model
model = CustomModel()

def predict(image: Image.Image):
    # סימולציה של זמן עיבוד
    time.sleep(2)  # שינוי הזמן לפי הצורך
    label, confidence = model.predict(image)
    result = "AI image" if label == 1 else "Real image"
    return result, f"Confidence: {confidence:.2f}%"

# Define the Gradio interface
with gr.Blocks() as demo:
    image_input = gr.Image(type="pil", label="Upload an image")
    animation = gr.Video("https://cdn-uploads.huggingface.co/production/uploads/66d6f1b3b50e35e1709bfdf7/x7Ud8PO9QPfmrTvBVcCKE.mp4", visible=False)
    output_label = gr.Textbox(label="Classification Result", interactive=False)
    output_confidence = gr.Textbox(label="Confidence", interactive=False)

    def show_animation(image):
        return animation.update(visible=True), "", ""  # להציג את האנימציה

    def hide_animation(image):
        result, confidence = predict(image)
        return animation.update(visible=False), result, confidence  # להסתיר את האנימציה ולהחזיר את התוצאות

    # עדכון על האירועים
    image_input.change(show_animation, inputs=image_input, outputs=[animation, output_label, output_confidence])
    image_input.submit(hide_animation, inputs=image_input, outputs=[output_label, output_confidence])

# Launch the Gradio interface
demo.launch()