Update app.py

app.py

@@ -1,43 +1,27 @@
 import gradio as gr
 import torch
-import torch.nn as nn
+import torch.nn.functional as F
 from torchvision import models, transforms
 from PIL import Image
-import json
-import os
 
-MODEL_PATH = "densenet121-a639ec97.pth"   # rename your .pth to model.pth
-
-# ------------------------------------------
-# Load ImageNet Label Mapping
-# ------------------------------------------
-with open("imagenet_classes.json", "r") as f:
-    idx_to_class = json.load(f)
+MODEL_PATH = "densenet121-a639ec97.pth"   # rename your .pth file to model.pth
 
 # ------------------------------------------
 # 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)
+    # Legacy checkpoints must use strict=False
     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
 
@@ -63,29 +47,28 @@ def predict(image):
     img_tensor = transform(image).unsqueeze(0)
 
     with torch.no_grad():
-        outputs = model(img_tensor)
-        probs = torch.softmax(outputs, dim=1)
+        logits = model(img_tensor)
+        probs = F.softmax(logits, 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": cls_name,
+            "class_index": int(idx.item()),
             "confidence": float(p.item())
         })
 
     return results
 
 # ------------------------------------------
-# Gradio Interface
+# Gradio UI
 # ------------------------------------------
-interface = gr.Interface(
+demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="numpy", label="Upload Image"),
-    outputs=gr.JSON(label="Top-5 Predictions"),
-    title="DenseNet-121 Legacy Model Classifier (ImageNet)",
-    description="Upload any image. Model returns top-5 ImageNet predictions.",
+    outputs=gr.JSON(label="Top-5 Predictions (Class Index Only)"),
+    title="DenseNet-121 Model Classifier (No Label File)",
+    description="This app uses a legacy DenseNet121 .pth model and returns top-5 class indices + confidence.",
 )
 
-interface.launch()
+demo.launch()