Update app.py

app.py

@@ -3,47 +3,76 @@ from transformers import ViTImageProcessor, ViTForImageClassification
 from PIL import Image
 import torch
 
-# --- CHANGEMENT CRITIQUE : Charger VOTRE modèle fine-tuné ---
-model_name = "MODLI/vit-fashion-classifier"  # <--- REMPLACER par votre modèle entraîné
+# --- Chargement du modèle et du processeur ---
+# Modèle de base ViT pré-entraîné sur ImageNet (beaucoup mieux que "beans")
+# C'est une solution temporaire en attendant de fine-tuner sur le dataset mode
+model_name = "google/vit-base-patch16-224"
 processor = ViTImageProcessor.from_pretrained(model_name)
 model = ViTForImageClassification.from_pretrained(model_name)
 
-# Fonction de prédiction avec seuil de confiance
 def predict(image):
-    # Pré-traiter l'image exactement comme pendant l'entraînement
-    inputs = processor(images=image, return_tensors="pt")
-    
-    # Prédire
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
-    
-    # Appliquer softmax pour obtenir les probabilités
-    probabilities = torch.nn.functional.softmax(logits, dim=-1)[0]
-    top_probs, top_indices = torch.topk(probabilities, 5)  # Top 5 predictions
-    
-    # Formater les résultats
-    predictions = []
-    for i, (prob, idx) in enumerate(zip(top_probs, top_indices)):
-        pred_label = model.config.id2label[idx.item()]
-        confidence = prob.item()
-        # N'afficher que si la confiance est > 5%
-        if confidence > 0.05 or i == 0:  # Toujours afficher la première même si faible
-            predictions.append((pred_label, f"{confidence:.2%}"))
-    
-    return predictions
+    """Fonction de prédiction avec gestion d'erreurs et seuil de confiance"""
+    try:
+        # Conversion vers RGB pour éviter les erreurs de canaux
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        # Pré-traitement de l'image
+        inputs = processor(images=image, return_tensors="pt")
+        
+        # Prédiction
+        with torch.no_grad():
+            outputs = model(**inputs)
+            logits = outputs.logits
+        
+        # Application de softmax pour obtenir les probabilités
+        probabilities = torch.nn.functional.softmax(logits, dim=-1)[0]
+        top_probs, top_indices = torch.topk(probabilities, 5)  # Top 5 predictions
+        
+        # Formatage des résultats
+        predictions = []
+        for i, (prob, idx) in enumerate(zip(top_probs, top_indices)):
+            pred_label = model.config.id2label[idx.item()]
+            confidence = prob.item()
+            # N'afficher que si la confiance est > 10%
+            if confidence > 0.1:
+                predictions.append(f"{pred_label}: {confidence:.2%}")
+        
+        # Si aucune prédiction n'a une confiance suffisante
+        if not predictions:
+            return "Je ne suis pas sûr de reconnaître cet item. Essayez avec une image plus claire."
+        
+        return "\n".join(predictions)
+        
+    except Exception as e:
+        return f"Une erreur s'est produite lors du traitement: {str(e)}"
 
-# Interface Gradio
+# Configuration de l'interface Gradio
 title = "Fashion Item Classifier"
-description = "Upload an image of a clothing item, and I will classify it."
+description = (
+    "Upload an image of a clothing item, and I will classify it. "
+    "⚠️ This is a general-purpose model. For better accuracy on fashion items, "
+    "a specialized model is needed."
+)
+
+# Exemples d'images (ajoutez vos propres exemples plus tard)
+examples = [
+    ["shirt_example.jpg"],
+    ["shoe_example.jpg"],
+    ["dress_example.jpg"]
+]
 
+# Création de l'interface
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil", label="Upload Clothing Item"),
-    outputs=gr.Label(label="Predictions", num_top_classes=5),
+    outputs=gr.Textbox(label="Classification Results"),
     title=title,
     description=description,
-    examples=[["path_to_example_image_1.jpg"], ["path_to_example_image_2.jpg"]], # Ajoutez des exemples
+    examples=examples,
+    allow_flagging="never"
 )
 
-demo.launch(debug=True)
+# Lancement de l'application
+if __name__ == "__main__":
+    demo.launch(debug=True, share=False)