Update app.py

app.py

@@ -7,42 +7,44 @@ import gradio as gr
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load class names (must match training order)
+# 1) Load class names from your saved file
 with open("classes.json", "r", encoding="utf-8") as f:
     class_names = json.load(f)
 
-# Build EfficientNet-B0 and replace classifier
+# 2) Build the model architecture (no downloading on the server)
 model = models.efficientnet_b0(weights=None)
 num_ftrs = model.classifier[1].in_features
 model.classifier[1] = nn.Linear(num_ftrs, len(class_names))
 
-# Load trained weights
+# 3) Load your trained weights
 state_dict = torch.load("best_efficientnetb0.pt", map_location=DEVICE)
 model.load_state_dict(state_dict)
 model.to(DEVICE)
 model.eval()
 
-# Same preprocessing as your val_transform
+# 4) Same preprocessing as validation/testing
 val_transform = Compose([
     Resize((224, 224)),
     ToTensor(),
-    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    Normalize(mean=[0.485, 0.456, 0.406],
+              std=[0.229, 0.224, 0.225]),
 ])
 
 @torch.no_grad()
 def predict(image):
-    # image comes as PIL
     x = val_transform(image).unsqueeze(0).to(DEVICE)
     logits = model(x)
     probs = torch.softmax(logits, dim=1)[0].cpu().tolist()
-    return {class_names[i]: float(probs[i]) for i in range(len(class_names))}
+    # show top-5 nicely
+    conf = {class_names[i]: float(probs[i]) for i in range(len(class_names))}
+    return conf
 
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil", label="Upload a dog image"),
     outputs=gr.Label(num_top_classes=5, label="Prediction"),
     title="Dog Breed Classifier (EfficientNet-B0)",
-    description="Upload a dog photo to get breed prediction + confidence.",
+    description="Upload a dog photo and the model predicts the breed with confidence."
 )
 
 if __name__ == "__main__":