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):
    print("Predict function called")  # הוספת שורת הדפסה
    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), "", ""

def show_results(image):
    result, confidence = predict(image)
    return gr.Video.update(visible=False), result, confidence

# יצירת ממשק Gradio עם אנימציות
with gr.Blocks(css="""
    .gr-button { 
        background-color: #4CAF50; 
        color: white; 
        transition: background-color 0.3s ease;
    }
    .gr-button:hover { 
        background-color: #45a049; 
    }
""") 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)
        
        output_label = gr.Textbox(label="Classification Result", interactive=False, visible=False)
        output_confidence = gr.Textbox(label="Confidence", interactive=False, visible=False)

    image_input.change(loading_animation, inputs=image_input, outputs=[animation, output_label, output_confidence])
    image_input.change(show_results, inputs=image_input, outputs=[animation, output_label, output_confidence])

    # הוספת כפתור
    submit_button = gr.Button("Submit")
    submit_button.click(predict, inputs=image_input, outputs=[output_label, output_confidence])

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