Update app.py

app.py

@@ -32,19 +32,16 @@ class ChineseClassifier(nn.Module):
         return x
 
 # -------- Utility Functions --------
-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 get_sorted_classes(labels_dict):
+    """Extract sorted unique classes from labels dictionary"""
+    return sorted(set(labels_dict.values()))
 
 def load_labels_json(labels_json_path):
-    """Load dict mapping image filename -> character label"""
+    """Load and normalize labels JSON"""
     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
+    # Normalize paths and remove directory prefixes
+    return {os.path.basename(k).replace("\\", "/"): v for k, v in labels_dict.items()}
 
 def prepare_transforms():
     return transforms.Compose([
@@ -70,35 +67,35 @@ def load_model(model_path, embed_dim, num_classes, device, pretrained=True, unfr
     model.eval()
     return model
 
-# -------- Globals and Setup --------
+# -------- Setup --------
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+EMBED_DIM = 512
+LABELS_JSON_PATH = "labels.json"
 
-# 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)
+# 1. Load labels and extract sorted classes
+labels_dict = load_labels_json(LABELS_JSON_PATH)
+classes = get_sorted_classes(labels_dict)
 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)
+# Verify class count matches training
+print(f"Loaded {num_classes} classes")
+print(f"First 5 classes: {classes[:5]}")
 
-# Download model weights from HF repo
+# 2. Download model
 REPO_ID = "JJJHHHH/CCR_EthicalSplit_Finetune"
 print("Downloading model from repo...")
 repo_dir = snapshot_download(repo_id=REPO_ID)
 model_path = os.path.join(repo_dir, "CCR_EthicalSplit_Finetune.pth")
-print("Model path:", model_path)
+print(f"Model path: {model_path}")
 
-# Load model and transforms
+# 3. Load model
 model = load_model(model_path, EMBED_DIM, num_classes, DEVICE)
 transform = prepare_transforms()
 
 # -------- Prediction Function --------
 def predict(pil_img):
+    """Predict character from PIL image"""
     img_t = transform(pil_img).unsqueeze(0).to(DEVICE)
     with torch.no_grad():
         output = model(img_t)
@@ -112,5 +109,5 @@ gr.Interface(
     inputs=gr.Image(type="pil", label="Upload Handwritten Chinese Character"),
     outputs=gr.Text(label="Predicted Character"),
     title="Chinese Character Recognition",
-    description="Upload an image of a handwritten Chinese character and get the predicted character."
-).launch(share=True)
+    description="Recognizes handwritten Chinese characters with 80% accuracy",
+).launch()