Update app.py

app.py

@@ -1,18 +1,17 @@
 import os
+import json
+from PIL import Image
 import torch
 import torch.nn as nn
-from PIL import Image
 from torchvision import models, transforms
 from huggingface_hub import snapshot_download
 import gradio as gr
 
-# --------------------------
-# Model Architecture
-# --------------------------
+# -------- Model Definition --------
 class ChineseClassifier(nn.Module):
     def __init__(self, embed_dim, num_classes, pretrainedEncoder=True, unfreezeEncoder=True):
         super().__init__()
-        resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT if pretrainedEncoder else None)
+        resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT) if pretrainedEncoder else models.resnet50()
         self.resnet = nn.Sequential(*list(resnet.children())[:-1])
         for param in self.resnet.parameters():
             param.requires_grad = unfreezeEncoder
@@ -29,14 +28,15 @@ class ChineseClassifier(nn.Module):
         x = self.dropout(x)
         if return_embedding:
             return x
-        return self.classifier(x)
+        x = self.classifier(x)
+        return x
 
-# --------------------------
-# Helper Functions
-# --------------------------
-def load_labels(path):
-    with open(path, "r", encoding="utf-8") as f:
-        return [line.strip() for line in f if line.strip()]
+# -------- Utility Functions --------
+def load_labels(labels_path):
+    # If your labels.txt is json-like, else adjust accordingly
+    with open(labels_path, "r", encoding="utf-8") as f:
+        labels = json.load(f)
+    return labels
 
 def prepare_transforms():
     return transforms.Compose([
@@ -46,41 +46,56 @@ def prepare_transforms():
                              std=[0.229, 0.224, 0.225]),
     ])
 
-def load_model(path, embed_dim, num_classes, device):
-    model = ChineseClassifier(embed_dim, num_classes).to(device)
-    model.load_state_dict(torch.load(path, map_location=device))
+def load_model(model_path, embed_dim, num_classes, device, pretrained=True, unfreeze=True):
+    model = ChineseClassifier(embed_dim, num_classes, pretrainedEncoder=pretrained, unfreezeEncoder=unfreeze).to(device)
+    checkpoint = torch.load(model_path, map_location=device)
+    if "model_state_dict" in checkpoint:
+        try:
+            model.load_state_dict(checkpoint["model_state_dict"])
+        except RuntimeError as e:
+            print("Warning:", e)
+            print("Loading partial weights, skipping classifier layer...")
+            filtered_state_dict = {k: v for k, v in checkpoint["model_state_dict"].items() if not k.startswith("classifier.")}
+            model.load_state_dict(filtered_state_dict, strict=False)
+    else:
+        model.load_state_dict(checkpoint)
     model.eval()
     return model
 
-# --------------------------
-# Init
-# --------------------------
+# -------- Globals and Setup --------
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load labels locally from Space repo root
+labels_path = "labels.txt"
+labels_dict = load_labels(labels_path)
+# Create list sorted by index (assuming labels_dict: filename->label)
+classes = sorted(set(labels_dict.values()))
+class_to_idx = {cls: idx for idx, cls in enumerate(classes)}
+idx_to_class = {v: k for k, v in class_to_idx.items()}
+num_classes = len(classes)
+EMBED_DIM = 512
+
+# Download model weights from HF repo
 REPO_ID = "JJJHHHH/CCR_EthicalSplit_Finetune"
+print("Downloading model from HF repo...")
 repo_dir = snapshot_download(repo_id=REPO_ID)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-labels_path = os.path.join(repo_dir, "labels.txt")
 model_path = os.path.join(repo_dir, "CCR_EthicalSplit_Finetune.pth")
+print("Model path:", model_path)
 
-class_names = load_labels(labels_path)
+# Prepare model and transforms
+model = load_model(model_path, EMBED_DIM, num_classes, DEVICE)
 transform = prepare_transforms()
-model = load_model(model_path, embed_dim=512, num_classes=len(class_names), device=device)
 
-# --------------------------
-# Inference
-# --------------------------
-def predict(image: Image.Image):
-    image = image.convert("RGB")
-    input_tensor = transform(image).unsqueeze(0).to(device)
+# -------- Prediction Function --------
+def predict(pil_img):
+    img_t = transform(pil_img).unsqueeze(0).to(DEVICE)
     with torch.no_grad():
-        output = model(input_tensor)
+        output = model(img_t)
         pred_idx = output.argmax(dim=1).item()
-        pred_label = class_names[pred_idx]
-    return f"Prediction: {pred_label}"
+        pred_label = idx_to_class[pred_idx]
+    return pred_label
 
-# --------------------------
-# Gradio UI
-# --------------------------
+# -------- Gradio Interface --------
 gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil", label="Upload Image"),