import gradio as gr
from transformers import AutoModelForImageClassification, AutoImageProcessor
from PIL import Image
import torch

# تحميل النموذج
model = AutoModelForImageClassification.from_pretrained(
    "asyafalni/arabichar-v3",
    trust_remote_code=True
)

# استخدام معالج صور جاهز
processor = AutoImageProcessor.from_pretrained(
    "google/vit-base-patch16-224-in21k"
)

# دالة لمعالجة الصورة واستخراج النص
def classify_image(image):
    try:
        # تحويل الصورة إلى RGB
        image = image.convert("RGB")

        # تجهيز المدخلات
        inputs = processor(images=image, return_tensors="pt")

        # تمريرها للنموذج
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            predicted_class = logits.argmax(-1).item()
            label = model.config.id2label.get(predicted_class, "Unknown")

        return label
    except Exception as e:
        return f"حدث خطأ: {e}"

# إنشاء واجهة Gradio
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil", label="رفع صورة"),
    outputs=gr.Textbox(label="النص المستخرج"),
    title="تصنيف الصور - Arabic Char",
    description="ارفع صورة وسيتم عرض النص المستخرج أسفلها."
)

# تشغيل الواجهة
iface.launch()