Update app.py

app.py

@@ -11,19 +11,28 @@ app = FastAPI()
 # Chargement des modèles
 def load_models():
     return {
-        "chest_classifier": pipeline("image-classification", model="codewithdark/vit-chest-xray")
+        "chest_classifier": pipeline("image-classification", model="codewithdark/vit-chest-xray", use_labels=True)
     }
 
 models = load_models()
 
-def translate_label(label):
-    translations = {
-        'Cardiomegaly': 'Kardiomegalie',
-        'Edema': 'Ödem',
-        'Consolidation': 'Konsolidierung',
-        'Pneumonia': 'Lungenentzündung',
-        'No Finding': 'Kein Befund'
-    }
+# 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
+]
+
+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})"
     return translations.get(label, label)
 
 def image_to_base64(image):
@@ -248,7 +257,7 @@ async def analyze_file(file: UploadFile = File(...)):
             results_html += f"""
                 <div>
                     <span class="{confidence_class}">{pred['score']:.1%}</span> - 
-                    {translate_label(pred['label'])}
+                    {translate_label(pred['label'] if isinstance(pred['label'], (int, str)) else pred['label'].item())}
                 </div>
             """