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

# Initialize the model
model = CustomModel()

def predict(image: Image.Image):
    # כאן מתבצע עיבוד התמונה
    label, confidence = model.predict(image)
    result = "AI image" if label == 1 else "Real image"
    return result, f"Confidence: {confidence:.2f}%"

# יצירת ממשק Gradio עם וידאו שמוצג בזמן חישוב המודל
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 as AI-generated or Real.",
    live=True,  # מאפשר להציג סטטוס בזמן אמת
    theme="compact",  # אפשרי להוסיף עיצוב מקצועי
    css=".loading-message { display: none; }"
)

# הגדרת הסרטון שיופיע בזמן חישוב
demo.loading = "https://cdn-uploads.huggingface.co/production/uploads/66d6f1b3b50e35e1709bfdf7/x7Ud8PO9QPfmrTvBVcCKE.mp4"  # הוסף את שם קובץ הווידאו

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