Update app.py

app.py

@@ -11,28 +11,26 @@ app = FastAPI()
 # Chargement des modèles
 def load_models():
     return {
-        "chest_classifier": pipeline("image-classification", model="codewithdark/vit-chest-xray", use_labels=True)
+        "chest_classifier": pipeline("image-classification", model="codewithdark/vit-chest-xray")
     }
 
 models = load_models()
 
 # Liste des maladies dans l'ordre des labels du modèle
-LABEL_MAP = [
-    'Kardiomegalie',  # Élargissement du cœur
-    'Ödem',           # Œdème
-    'Konsolidierung', # Consolidation
-    'Lungenentzündung', # Pneumonie
-    'Kein Befund'     # Aucune anomalie
-]
+LABEL_MAP = {
+    'LABEL_0': 'Kardiomegalie',  # Élargissement du cœur
+    'LABEL_1': 'Ödem',           # Œdème
+    'LABEL_2': 'Konsolidierung', # Consolidation
+    'LABEL_3': 'Lungenentzündung', # Pneumonie
+    'LABEL_4': 'Kein Befund'     # Aucune anomalie
+}
 
-def translate_label(index):
-    try:
-        # Convertir en entier si c'est une chaîne (ex: "LABEL_1" -> 1)
-        if isinstance(index, str) and index.startswith('LABEL_'):
-            index = int(index.split('_')[1]) - 1
-        return LABEL_MAP[index]
-    except (ValueError, IndexError):
-        return f"Unbekannt ({index})"
+def translate_label(label):
+    # Si le label est au format "LABEL_X", on le traduit
+    if isinstance(label, str) and label in LABEL_MAP:
+        return LABEL_MAP[label]
+    # Sinon on retourne le label original avec un avertissement
+    return f"Unbekannt ({label})"
     return translations.get(label, label)
 
 def image_to_base64(image):
@@ -257,7 +255,7 @@ async def analyze_file(file: UploadFile = File(...)):
             results_html += f"""
                 <div>
                     <span class="{confidence_class}">{pred['score']:.1%}</span> - 
-                    {translate_label(pred['label'] if isinstance(pred['label'], (int, str)) else pred['label'].item())}
+                    {translate_label(pred['label'])}
                 </div>
             """