Update app.py

app.py

@@ -4,29 +4,49 @@ import torch.nn as nn
 from torchvision import models, transforms
 from PIL import Image
 import json
+import os
 
-# --------------------------
-# Load the model
-# --------------------------
-model_path = "densenet121-a639ec97.pth"  # put this file in the Space
+MODEL_PATH = "densenet121-a639ec97.pth"   # rename your .pth to model.pth
 
-state_dict = torch.load(model_path, map_location="cpu")
-
-model = models.densenet121(weights=None)
-model.load_state_dict(state_dict, strict=True)
-model.eval()
-
-# --------------------------
-# Load ImageNet class labels
-# --------------------------
-# Downloaded from: https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt
+# ------------------------------------------
+# Load ImageNet Label Mapping
+# ------------------------------------------
 with open("imagenet_classes.json", "r") as f:
     idx_to_class = json.load(f)
 
-# --------------------------
-# Preprocessing function
-# --------------------------
-preprocess = transforms.Compose([
+# ------------------------------------------
+# Load Model with Legacy-Key Compatibility
+# ------------------------------------------
+def load_legacy_densenet(path):
+    print("Loading model:", path)
+
+    # Load state_dict (old torchvision format)
+    state_dict = torch.load(path, map_location="cpu")
+
+    # Create new DenseNet model
+    model = models.densenet121(weights=None)
+
+    # Try loading with strict=False (legacy fix)
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+
+    print("\n=== LOADING SUMMARY ===")
+    print("Missing keys:", len(missing))
+    print("Unexpected keys:", len(unexpected))
+
+    if len(missing) > 0:
+        print("⚠ NOTE: Missing keys detected (normal for legacy checkpoint)")
+    if len(unexpected) > 0:
+        print("⚠ NOTE: Unexpected keys detected (normal for legacy checkpoint)")
+
+    model.eval()
+    return model
+
+model = load_legacy_densenet(MODEL_PATH)
+
+# ------------------------------------------
+# Preprocessing
+# ------------------------------------------
+transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
     transforms.Normalize(
@@ -35,36 +55,37 @@ preprocess = transforms.Compose([
     )
 ])
 
-# --------------------------
+# ------------------------------------------
 # Prediction function
-# --------------------------
+# ------------------------------------------
 def predict(image):
     image = Image.fromarray(image).convert("RGB")
-    img_tensor = preprocess(image).unsqueeze(0)
+    img_tensor = transform(image).unsqueeze(0)
 
     with torch.no_grad():
         outputs = model(img_tensor)
-        probs = torch.nn.functional.softmax(outputs, dim=1)
+        probs = torch.softmax(outputs, dim=1)
         top5_prob, top5_idx = torch.topk(probs, 5)
 
     results = []
     for p, idx in zip(top5_prob[0], top5_idx[0]):
+        cls_name = idx_to_class.get(str(idx.item()), "Unknown")
         results.append({
-            "class": idx_to_class[str(idx.item())],
+            "class": cls_name,
             "confidence": float(p.item())
         })
 
     return results
 
-# --------------------------
-# Gradio UI
-# --------------------------
-demo = gr.Interface(
+# ------------------------------------------
+# Gradio Interface
+# ------------------------------------------
+interface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="numpy", label="Upload Image"),
     outputs=gr.JSON(label="Top-5 Predictions"),
-    title="DenseNet-121 ImageNet Classifier",
-    description="Upload an image and get top-5 ImageNet predictions.",
+    title="DenseNet-121 Legacy Model Classifier (ImageNet)",
+    description="Upload any image. Model returns top-5 ImageNet predictions.",
 )
 
-demo.launch()
+interface.launch()