Update app.py

app.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI, Query
 from transformers import AutoModelForImageClassification, AutoImageProcessor
-from PIL import Image
+from PIL import Image, UnidentifiedImageError
 import requests
 from io import BytesIO
 import torch
@@ -30,9 +30,15 @@ def root():
 @app.post("/classify")
 def classify_image(url: str = Query(..., description="رابط الصورة")):
     try:
-        # تحميل الصورة
-        response = requests.get(url)
-        image = Image.open(BytesIO(response.content)).convert("RGB")
+        # تحميل الصورة مع timeout ومع التحقق
+        response = requests.get(url, timeout=10, stream=True)
+        response.raise_for_status()  # رفع استثناء لو كان الرابط غير صالح
+
+        # قراءة الصورة
+        try:
+            image = Image.open(BytesIO(response.content)).convert("RGB")
+        except UnidentifiedImageError:
+            return {"error": "الرابط لا يحتوي على صورة صالحة"}
 
         # تجهيز المدخلات
         inputs = processor(images=image, return_tensors="pt")
@@ -45,9 +51,11 @@ def classify_image(url: str = Query(..., description="رابط الصورة")):
 
         return {
             "url": url,
-            "predicted_class": model.config.id2label[predicted_class],
+            "predicted_class": model.config.id2label.get(predicted_class, "Unknown"),
             "scores": logits.softmax(-1).tolist()
         }
 
+    except requests.exceptions.RequestException as e:
+        return {"error": f"تعذر الوصول إلى الرابط: {e}"}
     except Exception as e:
         return {"error": str(e)}