Add application file

app.py

@@ -2,27 +2,53 @@ import gradio as gr
 from transformers import AutoFeatureExtractor, AutoModelForImageClassification
 from PIL import Image
 import torch
+import torch.nn.functional as F
+import pandas as pd
 
-extractor = AutoFeatureExtractor.from_pretrained("microsoft/resnet-50")
-model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-50")
+extractor = AutoFeatureExtractor.from_pretrained("nateraw/resnet50-trash-classifier")
+model = AutoModelForImageClassification.from_pretrained("nateraw/resnet50-trash-classifier")
 
 POUBELLES = {
-    "banana": "biodéchets",
-    "plastic bottle": "plastique",
-    "can": "métal",
-    "apple": "biodéchets",
-    "paper towel": "papier",
+    "plastic": "plastique",
     "glass": "verre",
+    "metal": "métal",
+    "paper": "papier",
+    "cardboard": "papier/carton",
+    "trash": "ordures ménagères",
+    "compost": "biodéchets",
+    "battery": "déchet dangereux",
+    "clothes": "textile",
+    # Ajoute d'autres si nécessaire
 }
 
 def classify_image(image):
     inputs = extractor(images=image, return_tensors="pt")
     with torch.no_grad():
         logits = model(**inputs).logits
-    predicted_class_idx = logits.argmax(-1).item()
-    label = model.config.id2label[predicted_class_idx]
+        probs = F.softmax(logits, dim=-1)
 
-    poubelle = POUBELLES.get(label.lower(), "inconnue")
-    return f"{label} → {poubelle}"
+    top_probs, top_idxs = torch.topk(probs, 3)
+    top_probs = top_probs.squeeze().tolist()
+    top_idxs = top_idxs.squeeze().tolist()
 
-gr.Interface(fn=classify_image, inputs=gr.Image(type="pil"), outputs="text").launch()
+    rows = []
+    for idx, prob in zip(top_idxs, top_probs):
+        label = model.config.id2label[idx]
+        poubelle = POUBELLES.get(label.lower(), "inconnue")
+        rows.append({
+            "Objet": label,
+            "Poubelle": poubelle,
+            "Confiance (%)": round(prob * 100, 2),
+        })
+
+    return pd.DataFrame(rows)
+
+interface = gr.Interface(
+    fn=classify_image,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Dataframe(),
+    title="🗑️ Classifieur de Déchets",
+    description="Dépose une image de déchet pour savoir dans quelle poubelle le trier."
+)
+
+interface.launch()