Update app.py

app.py

@@ -1,7 +1,7 @@
 import json, os
 import numpy as np
 from PIL import Image
-import ast
+
 
 
 import torch
@@ -10,6 +10,31 @@ import timm
 from timm.data import resolve_model_data_config, create_transform
 from transformers import AutoTokenizer, AutoModel
 import gradio as gr
+import ast
+from huggingface_hub import hf_hub_download
+
+SPACE_REPO = os.getenv("SPACE_REPO_NAME", "muruga778/api_for_model")  # change if your space id differs
+
+def safe_torch_load(filename: str):
+    """
+    1) try local file
+    2) if corrupted -> force-download from Hub cache and load again
+    """
+    try:
+        print(f"🔎 Loading weights: {filename} (local)")
+        return torch.load(filename, map_location="cpu")
+    except Exception as e:
+        print(f"⚠️ Local load failed for {filename}: {repr(e)}")
+        print("⬇️ Force-downloading from Hugging Face Hub cache...")
+        cached = hf_hub_download(
+            repo_id=SPACE_REPO,
+            repo_type="space",
+            filename=filename,
+            force_download=True,
+        )
+        print("✅ Downloaded to:", cached, "size(MB)=", os.path.getsize(cached)/1024/1024)
+        return torch.load(cached, map_location="cpu")
+
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -59,35 +84,26 @@ def triage(label, conf, text):
 CFG = load_json("fusion_config.json")
 LABEL_MAP = load_json("label_map.json")
 
-# ✅ your label_map.json is a wrapper: {"classes":[...], "label2idx":{...}}
-if isinstance(LABEL_MAP, dict) and "classes" in LABEL_MAP and "label2idx" in LABEL_MAP:
-    CLASSES = [str(x).strip().lower() for x in LABEL_MAP["classes"]]
-    label2idx = {str(k).strip().lower(): int(v) for k, v in LABEL_MAP["label2idx"].items()}
-
-    # (optional but good) sanity check: classes order matches label2idx
-    if len(CLASSES) != len(label2idx):
-        raise ValueError("label_map.json mismatch: len(classes) != len(label2idx)")
-    for i, lab in enumerate(CLASSES):
-        if label2idx.get(lab, None) != i:
-            # if mismatch, rebuild CLASSES from label2idx index order
-            idx2label = {int(v): str(k).strip().lower() for k, v in LABEL_MAP["label2idx"].items()}
-            CLASSES = [idx2label[i] for i in sorted(idx2label.keys())]
-            break
+# Your label_map.json looks like: {"classes":[...], "label2idx":{...}}
+if isinstance(LABEL_MAP, dict) and "classes" in LABEL_MAP and isinstance(LABEL_MAP["classes"], list):
+    CLASSES = [str(x) for x in LABEL_MAP["classes"]]
+    label2idx = LABEL_MAP.get("label2idx", {c: i for i, c in enumerate(CLASSES)})
+
+# Older possible formats:
+elif isinstance(LABEL_MAP, dict) and all(isinstance(k, str) and k.isdigit() for k in LABEL_MAP.keys()):
+    # {"0":"eczema", ...}
+    idx2label = {int(k): str(v) for k, v in LABEL_MAP.items()}
+    CLASSES = [idx2label[i] for i in sorted(idx2label.keys())]
+    label2idx = {c: i for i, c in enumerate(CLASSES)}
 
 else:
-    # Fallback: support plain formats too
-    if all(isinstance(k, str) and k.isdigit() for k in LABEL_MAP.keys()):  # {"0":"eczema",...}
-        idx2label = {int(k): str(v).strip().lower() for k, v in LABEL_MAP.items()}
-        CLASSES = [idx2label[i] for i in sorted(idx2label.keys())]
-        label2idx = {c: i for i, c in enumerate(CLASSES)}
-    else:  # {"eczema":0,...}
-        label2idx = {str(k).strip().lower(): int(v) for k, v in LABEL_MAP.items()}
-        CLASSES = [c for c, _ in sorted(label2idx.items(), key=lambda x: x[1])]
+    # {"eczema": 0, ...}
+    label2idx = {str(k): int(v) for k, v in LABEL_MAP.items()}
+    CLASSES = [c for c, _ in sorted(label2idx.items(), key=lambda x: x[1])]
 
 NUM_CLASSES = len(CLASSES)
-
-print("✅ NUM_CLASSES:", NUM_CLASSES)          # should print 16
-print("✅ First labels:", CLASSES[:5])         # sanity check
+print("✅ NUM_CLASSES:", NUM_CLASSES)
+print("✅ First labels:", CLASSES[:5])
 
 IMG_BACKBONE = CFG.get("img_backbone", "tf_efficientnetv2_s")
 IMG_SIZE = int(CFG.get("img_size", 384))