test

app.py

@@ -1,15 +1,33 @@
 import gradio as gr
-from transformers import AutoImageProcessor, AutoModelForImageClassification
-from PIL import Image
+from transformers import pipeline, AutoImageProcessor, AutoModelForImageClassification
+from PIL import Image, ImageDraw
 import torch
-import torch.nn.functional as F
+from torchvision import transforms
 import pandas as pd
 
-# Chargement du processeur et du modèle
-processor = AutoImageProcessor.from_pretrained("tribber93/my-trash-classification")
-model = AutoModelForImageClassification.from_pretrained("tribber93/my-trash-classification")
+# -------------------------------------
+# 1. DATA AUGMENTATION POUR ROBUSTESSE
+# -------------------------------------
+augment = transforms.Compose([
+    transforms.RandomHorizontalFlip(p=0.5),
+    transforms.RandomRotation(10),
+    transforms.ColorJitter(brightness=0.2, contrast=0.2),
+])
 
-# Dictionnaire de correspondance entre les labels et les types de poubelles
+# -------------------------------------
+# 2. CHARGEMENT DU MODELE FINE-TUNED TRASHNET
+# -------------------------------------
+MODEL_ID = "mrm8488/resnet50-finetuned-trashnet"
+trash_classifier = pipeline(
+    "image-classification",
+    model=MODEL_ID,
+    device=0 if torch.cuda.is_available() else -1,
+    top_k=3
+)
+
+# -------------------------------------
+# 3. DICTIONNAIRE DE MAPPING
+# -------------------------------------
 POUBELLES = {
     "cardboard": "papier/carton",
     "glass": "verre",
@@ -19,34 +37,49 @@ POUBELLES = {
     "trash": "ordures ménagères",
 }
 
-# Fonction de classification de l'image
-def classify_image(image):
-    inputs = processor(images=image, return_tensors="pt")
-    with torch.no_grad():
-        logits = model(**inputs).logits
-        probs = F.softmax(logits, dim=-1)
+# -------------------------------------
+# 4. FONCTION DE CLASSIFICATION
+# -------------------------------------
+def classify_image(image: Image.Image):
+    # Appliquer data augmentation pour robustesse
+    image_aug = augment(image)
 
-    top_probs, top_idxs = torch.topk(probs, 3)
-    top_probs = top_probs.squeeze().tolist()
-    top_idxs = top_idxs.squeeze().tolist()
+    # Prédiction via pipeline
+    results = trash_classifier(image_aug)
 
+    # Préparer les DataFrame rows
     rows = []
-    for idx, prob in zip(top_idxs, top_probs):
-        label = model.config.id2label[idx]
+    for r in results:
+        label = r["label"]
+        score = r["score"]
         poubelle = POUBELLES.get(label.lower(), "inconnue")
         rows.append({
             "Objet": label,
             "Poubelle": poubelle,
-            "Confiance (%)": round(prob * 100, 2),
+            "Confiance (%)": round(score * 100, 2)
         })
 
+    # Renvoie un tableau pandas pour Gradio
     return pd.DataFrame(rows)
 
-# Création de l'interface Gradio
-gr.Interface(
+# -------------------------------------
+# 5. LANCEMENT DE L'INTERFACE GRADIO
+# -------------------------------------
+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."
-).launch()
+    outputs=gr.Dataframe(
+        headers=["Objet", "Poubelle", "Confiance (%)"],
+        row_count=(1, 10)
+    ),
+    title="🗑️ Classifieur de Déchets Amélioré",
+    description=(
+        "Dépose une image de déchet pour savoir dans quelle poubelle la trier. "
+        "Le modèle est fine-tuné sur TrashNet et bénéficie de data augmentation pour une meilleure robustesse."
+    ),
+    examples=None,
+    allow_flagging="never"
+)
+
+if __name__ == "__main__":
+    interface.launch()