Update app.py

app.py

@@ -13,7 +13,7 @@ from pymongo import MongoClient
 import warnings
 
 # =====================
-# Silence HF warnings (safe)
+# Silence HF warnings
 # =====================
 warnings.filterwarnings("ignore")
 
@@ -35,7 +35,7 @@ db = client["skin-disease-db"]
 reports = db["reports"]
 
 # =====================
-# Labels
+# Labels (ORDER MUST MATCH TRAINING)
 # =====================
 labels = [
     "Acne and Rosacea Photos",
@@ -67,40 +67,36 @@ NUM_CLASSES = len(labels)
 device = torch.device("cpu")
 
 # =====================
-# Load model weights from HF Hub
+# Load trained model
 # =====================
 weights_path = hf_hub_download(
     repo_id="pragun3669/dermify-vit",
     filename="best_vit1_model.pth"
 )
 
-# Base ViT architecture (23 classes)
 model = ViTForImageClassification.from_pretrained(
     "google/vit-large-patch16-224",
     num_labels=NUM_CLASSES,
     ignore_mismatched_sizes=True
 )
 
-# Load checkpoint safely (IGNORE classifier mismatch)
+# ✅ LOAD FULL TRAINED STATE (INCLUDING CLASSIFIER)
 state_dict = torch.load(weights_path, map_location=device)
-
-# 🔥 Drop classifier weights if they exist (1000-class head)
-state_dict = {
-    k: v for k, v in state_dict.items()
-    if not k.startswith("classifier")
-}
-
 model.load_state_dict(state_dict, strict=False)
+
 model.to(device)
 model.eval()
 
 # =====================
-# Image Transform
+# Image Transform (MATCH TRAINING)
 # =====================
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+    transforms.Normalize(
+        mean=[0.5, 0.5, 0.5],
+        std=[0.5, 0.5, 0.5]
+    )
 ])
 
 # =====================
@@ -126,7 +122,7 @@ def predict():
         logits = model(tensor).logits
         probs = F.softmax(logits, dim=1)
 
-    idx = probs.argmax().item()
+    idx = probs.argmax(dim=1).item()
     confidence = probs[0][idx].item()
 
     reports.insert_one({
@@ -141,7 +137,7 @@ def predict():
     })
 
 # =====================
-# Run (HF Spaces compatible)
+# Run (HF Spaces)
 # =====================
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)