Update app.py

app.py

@@ -1,6 +1,8 @@
 import json, os
 import numpy as np
 from PIL import Image
+import ast
+
 
 import torch
 import torch.nn as nn
@@ -58,15 +60,67 @@ CFG = load_json("fusion_config.json")
 LABEL_MAP = load_json("label_map.json")
 
 # label_map can be {"label": idx} or {"0":"label"}
-if all(isinstance(k, str) and k.isdigit() for k in LABEL_MAP.keys()):
-    idx2label = {int(k): v for k, v in LABEL_MAP.items()}
-    CLASSES = [idx2label[i] for i in sorted(idx2label.keys())]
-else:
-    label2idx = {k: int(v) for k, v in LABEL_MAP.items()}
-    CLASSES = [c for c, _ in sorted(label2idx.items(), key=lambda x: x[1])]
 
+def normalize_label(x):
+    """
+    Makes label_map robust:
+    - ["eczema"] -> "eczema"
+    - "['eczema']" -> "eczema"
+    - [] / "" / None -> "unknown"
+    """
+    if x is None:
+        return "unknown"
+    if isinstance(x, list):
+        return normalize_label(x[0]) if len(x) else "unknown"
+    s = str(x).strip()
+    if s == "" or s.lower() in ["none", "nan", "null"]:
+        return "unknown"
+    if s.startswith("[") and s.endswith("]"):
+        try:
+            return normalize_label(ast.literal_eval(s))
+        except Exception:
+            pass
+    return s.lower()
+
+def build_classes(label_map):
+    if not isinstance(label_map, dict) or len(label_map) == 0:
+        raise ValueError("label_map.json must be a non-empty dict")
+
+    keys = list(label_map.keys())
+    vals = list(label_map.values())
+
+    # Case A: {"0":"eczema", "1":"acne"}  (idx -> label)
+    if all(isinstance(k, str) and k.isdigit() for k in keys):
+        idx2label = {int(k): normalize_label(v) for k, v in label_map.items()}
+        classes = [idx2label[i] for i in sorted(idx2label.keys())]
+        return classes
+
+    # Case B: {"eczema":0, "acne":1}  (label -> idx)
+    if all(isinstance(v, (int, float)) and float(v).is_integer() for v in vals):
+        label2idx = {normalize_label(k): int(v) for k, v in label_map.items()}
+        classes = [lab for lab, _ in sorted(label2idx.items(), key=lambda x: x[1])]
+        return classes
+
+    # Case C: weird format like {"eczema": ["eczema"], ...}
+    # Use KEYS as labels (dedup + sorted)
+    classes = sorted({normalize_label(k) for k in keys})
+    return classes
+
+CLASSES = build_classes(LABEL_MAP)
 NUM_CLASSES = len(CLASSES)
 
+print("✅ NUM_CLASSES:", NUM_CLASSES)
+print("✅ First labels:", CLASSES[:10])
+
+IMG_BACKBONE = CFG.get("img_backbone", "tf_efficientnetv2_s")
+IMG_SIZE = int(CFG.get("img_size", 384))
+TEXT_MODEL_NAME = CFG.get(
+    "text_model_name",
+    "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext"
+)
+MAX_LEN = int(CFG.get("max_len", 128))
+
+
 IMG_BACKBONE = CFG.get("img_backbone", "tf_efficientnetv2_s")
 IMG_SIZE = int(CFG.get("img_size", 384))
 TEXT_MODEL_NAME = CFG.get("text_model_name", "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext")