import gradio as gr
from PIL import Image, UnidentifiedImageError
import numpy as np
import requests
from io import BytesIO
import json

from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline

# معرف النموذج الخاص بك على Hugging Face
model_id = "smartgmin/Entrnal_5class_agumm_last_newV7_model"

# تحميل النموذج ومعالج الميزات مرة واحدة عند بدء التطبيق
model2 = AutoModelForImageClassification.from_pretrained(model_id, from_tf=True)
feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
clsss = pipeline('image-classification', model=model2, feature_extractor=feature_extractor)

def predict(image_url):
    try:
        # التحقق من صحة الرابط وتحميل الصورة
        response = requests.get(image_url, timeout=10)
        response.raise_for_status()  # سيقوم برفع استثناء إذا كان الطلب غير ناجح

        # التحقق من نوع المحتوى
        content_type = response.headers.get('Content-Type')
        if not content_type or not content_type.startswith('image'):
            return json.dumps({"error": "الرابط المقدم لا يشير إلى صورة صالحة."})

        # تحميل الصورة من الرابط
        image = Image.open(BytesIO(response.content)).convert("RGB")

        # تمرير الصورة إلى النموذج
        yl = clsss(image)
        max_item = max(yl, key=lambda x: x['score'])
        nn = "{:.2f}".format(max_item['score'])  # تنسيق الدقة ليكون مقروءًا بشكل أفضل
        dd = max_item['label']
        
        # إرجاع النتيجة بصيغة JSON
        return json.dumps({"label": dd, "score": nn})
    
    except requests.exceptions.RequestException as e:
        return json.dumps({"error": f"خطأ في تحميل الصورة من الرابط المقدم: {e}"})
    except UnidentifiedImageError:
        return json.dumps({"error": "لا يمكن تحديد نوع الصورة من الرابط المقدم."})
    except Exception as e:
        return json.dumps({"error": f"حدث خطأ أثناء معالجة الصورة: {e}"})

# إنشاء واجهة Gradio باستخدام المكونات الجديدة
iface = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=2, placeholder="أدخل رابط الصورة هنا...", label="رابط الصورة"),
    outputs="text",  # مخرجات النص تُستخدم هنا لعرض نص الـ JSON
    title="نموذج ViT لتصنيف الصور",
    description="أدخل رابط صورة للحصول على تصنيف باستخدام نموذج ViT المدرب."
)

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