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(5)  # ניתן לשנות את זמן העיבוד לסימולציה
    
    # הפעלת המודל כדי לסווג את התמונה
    label, confidence = model.predict(image)
    result = "AI image" if label == 1 else "Real image"
    
    return result, f"Confidence: {confidence:.2f}%"

def loading_animation(image):
    # מציג וידאו עד שתוצאת המודל מתקבלת
    return gr.Video.update(visible=True), gr.Textbox.update(visible=False), gr.Textbox.update(visible=False)

def show_results(image):
    # מפעיל את המודל ומחזיר את התוצאות
    result, confidence = predict(image)
    return gr.Video.update(visible=False), result, confidence

# יצירת ממשק Gradio
with gr.Blocks() as demo:
    with gr.Row():
        # קלט תמונה
        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, elem_id="loading_video")
        
        # תוצאות
        output_label = gr.Textbox(label="Classification Result", visible=False)
        output_confidence = gr.Textbox(label="Confidence", visible=False)

    # אירועים
    image_input.change(loading_animation, inputs=image_input, outputs=[animation, output_label, output_confidence])
    image_input.submit(show_results, inputs=image_input, outputs=[animation, output_label, output_confidence])

# השקת הממשק
demo.launch()