Update app.py

app.py

@@ -32,10 +32,19 @@ class ChineseClassifier(nn.Module):
         return x
 
 # -------- Utility Functions --------
-def load_labels(labels_path):
-    with open(labels_path, "r", encoding="utf-8") as f:
-        labels = [line.strip() for line in f if line.strip()]
-    return labels
+def load_class_list(labels_txt_path):
+    """Load ordered list of classes (characters) from labels.txt"""
+    with open(labels_txt_path, "r", encoding="utf-8") as f:
+        classes = [line.strip() for line in f if line.strip()]
+    return classes
+
+def load_labels_json(labels_json_path):
+    """Load dict mapping image filename -> character label"""
+    with open(labels_json_path, "r", encoding="utf-8") as f:
+        labels_dict = json.load(f)
+    # Normalize Windows-style backslash paths to slash paths
+    labels_dict = {k.replace("\\", "/"): v for k, v in labels_dict.items()}
+    return labels_dict
 
 def prepare_transforms():
     return transforms.Compose([
@@ -64,13 +73,19 @@ def load_model(model_path, embed_dim, num_classes, device, pretrained=True, unfr
 # -------- Globals and Setup --------
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load labels from JSON
-labels_path = "labels.txt"
-labels = load_labels(labels_path)  # List of characters ordered by index
-idx_to_class = {idx: label for idx, label in enumerate(labels)}
-num_classes = len(labels)
+# Paths to your label files (make sure these are accessible in your environment)
+LABELS_TXT_PATH = "labels.txt"        # Your class list: idx -> character
+LABELS_JSON_PATH = "labels.json"      # Your filename -> character mapping (optional, for evaluation)
+
+# Load class list for prediction indexing
+classes = load_class_list(LABELS_TXT_PATH)
+idx_to_class = {idx: c for idx, c in enumerate(classes)}
+num_classes = len(classes)
 EMBED_DIM = 512
 
+# Load the labels.json if you want (not required for prediction)
+# filename_to_char = load_labels_json(LABELS_JSON_PATH)
+
 # Download model weights from HF repo
 REPO_ID = "JJJHHHH/CCR_EthicalSplit_Finetune"
 print("Downloading model from repo...")