Update app.py

app.py

@@ -23,22 +23,29 @@ id2label = {
 }
 
 def predict(image):
-    # Preprocess the image
+    # Prétraiter l'image
     inputs = feature_extractor(images=image, return_tensors="pt")
     
-    # Forward pass through the model
+    # Passage en avant dans le modèle
     outputs = model(**inputs)
     
-    # Get the predicted label and confidence score
+    # Obtenir les logits
     logits = outputs.logits
-    predicted_class_idx = logits.argmax(-1).item()
-    confidence_score = torch.nn.functional.softmax(logits, dim=-1)[0, predicted_class_idx].item()
     
-    # Get the label name
-    predicted_label = id2label[str(predicted_class_idx)]
+    # Calculer les scores de confiance avec softmax
+    probs = torch.nn.functional.softmax(logits, dim=-1)[0]
     
-    # Return the predicted label and confidence score
-    return f"{predicted_label}: {confidence_score:.2f}"
+    # Obtenir les indices des trois classes les plus probables
+    top_3_indices = torch.topk(probs, 3).indices.tolist()
+    
+    # Obtenir les labels et les scores de confiance pour les trois classes les plus probables
+    top_3_labels_and_scores = [(id2label[str(idx)], probs[idx].item()) for idx in top_3_indices]
+    
+    # Formater les résultats
+    results = [f"{label}: {score:.2f}" for label, score in top_3_labels_and_scores]
+    
+    # Retourner les résultats
+    return "\n".join(results)
 
     
 # Create the Gradio interface