Update app.py

app.py

@@ -1,44 +1,4 @@
-Hugging Face's logo
-Hugging Face
-Models
-Datasets
-Spaces
-Docs
-Pricing
-
-
-Spaces:
-NiranjanIITPKD
-/
-tinyimagenet-resnet18-demo 
-
-like
-0
-
-App
-Files
-Community
-Settings
-tinyimagenet-resnet18-demo
-/
-app.py
-
-NiranjanIITPKD's picture
-NiranjanIITPKD
-Create app.py
-5818cdb
-verified
-1 day ago
-raw
-
-Copy download link
-history
-blame
-edit
-delete
-
-1.3 kB
-import torch
+mport torch
 import torch.nn as nn
 from torchvision import models, transforms as T
 from PIL import Image
@@ -46,18 +6,21 @@ import gradio as gr
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load label mappings
+# Load label files
 wnids = [line.strip() for line in open("wnids.txt")]
 
-# Map wnid → human class name
+# Map wnid → readable label
 words = {}
 with open("words.txt", "r") as f:
     for line in f:
         wnid, name = line.split("\t")
-        words[wnid] = name.split(",")[0]  # take first name only
+        words[wnid] = name.split(",")[0]
+
+# ✅ CORRECT CLASS ORDER (alphabetical)
+sorted_wnids = sorted(wnids)
 
-# Create final label list (human readable)
-id_to_label = [words[wnid] for wnid in wnids]
+# ✅ Final label list matching model training order
+id_to_label = [words[wnid] for wnid in sorted_wnids]
 
 # Load Model
 model = models.resnet18(weights=None)
@@ -66,7 +29,7 @@ model.load_state_dict(torch.load("best_resnet18_tinyimagenet.pth", map_location=
 model.to(device)
 model.eval()
 
-# Same transform used in training
+# Same preprocessing as training
 transform = T.Compose([
     T.Resize((224, 224)),
     T.ToTensor(),
@@ -85,5 +48,5 @@ gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title="Tiny ImageNet Classifier (ResNet18)"
+    title="Tiny ImageNet Classifier (Corrected Labels)"
 ).launch()