Update app.py

app.py

@@ -3,37 +3,73 @@ import onnxruntime as ort
 import numpy as np
 from PIL import Image
 from huggingface_hub import hf_hub_download
+import base64
+import io
 
-model_path = hf_hub_download(repo_id="cabrel09/insect-detection-model", filename="vit_insects.onnx")
-hf_hub_download(repo_id="cabrel09/insect-detection-model", filename="vit_insects.onnx.data")
-session = ort.InferenceSession(model_path)
-
-# L'ordre crucial des classes
+# Labels exacts
 labels = [
     "Fall Armyworms", "Western Corn Rootworms", "Colorado Potato Beetles", "Thrips", 
     "Corn Earworms", "Cabbage Loopers", "Armyworms", "Brown Marmorated Stink Bugs", 
     "Tomato Hornworms", "Citrus Canker", "Aphids", "Corn Borers", "Fruit Flies", 
-    "Africanized Honey Bees", "Spider Mites"
+    "Africanized Honey Bees (Killer Bees)", "Spider Mites"
 ]
 
-def predict(img):
+# Chargement sécurisé
+try:
+    model_path = hf_hub_download(repo_id="cabrel09/insect-detection-model", filename="vit_insects.onnx")
+    hf_hub_download(repo_id="cabrel09/insect-detection-model", filename="vit_insects.onnx.data")
+    session = ort.InferenceSession(model_path)
+except Exception as e:
+    print(f"Erreur chargement: {e}")
+    session = None
+
+def predict(img_input):
+    if session is None: return {"error": "Modèle non chargé"}
+    if img_input is None: return None
+        
     try:
+        # GESTION HYBRIDE : Détection automatique du type d'entrée
+        if isinstance(img_input, str):
+            # Si c'est du base64 (via API)
+            if "base64," in img_input:
+                img_input = img_input.split("base64,")[1]
+            img_bytes = base64.b64decode(img_input)
+            img = Image.open(io.BytesIO(img_bytes))
+        elif isinstance(img_input, dict) and "url" in img_input:
+            # Si Gradio envoie un dictionnaire (via API structurée)
+            url_data = img_input["url"]
+            if "base64," in url_data:
+                url_data = url_data.split("base64,")[1]
+            img_bytes = base64.b64decode(url_data)
+            img = Image.open(io.BytesIO(img_bytes))
+        else:
+            # Si c'est déjà une image PIL (via l'interface Web)
+            img = img_input
+
+        # Preprocessing
         img = img.convert("RGB").resize((224, 224))
         img_array = np.array(img).astype('float32') / 255.0
         img_array = (img_array - 0.5) / 0.5
         img_array = np.transpose(img_array, (2, 0, 1))
         img_array = np.expand_dims(img_array, axis=0)
 
+        # Inférence
         outputs = session.run(None, {session.get_inputs()[0].name: img_array})
         logits = outputs[0][0]
-        
         probs = np.exp(logits - np.max(logits))
         probs /= probs.sum()
 
-        # On associe chaque probabilité à son nom d'insecte
-        return {labels[i]: float(probs[i]) for i in range(len(labels))}
+        return {labels[i]: float(probs[i]) for i in range(min(len(labels), len(probs)))}
+        
     except Exception as e:
         return {"error": str(e)}
 
-interface = gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs=gr.Label(num_top_classes=5))
-interface.launch()
+# Utilisation de gr.Image mais acceptation de types flexibles
+demo = gr.Interface(
+    fn=predict, 
+    inputs=gr.Image(type="pil"), 
+    outputs=gr.Label(num_top_classes=5),
+    title="PlantPatrol Classifier"
+)
+
+demo.launch()