Update app.py

app.py

@@ -4,13 +4,28 @@ import torch.nn as nn
 import torchvision.transforms as T
 import torchvision.models as models
 from PIL import Image
-import json
 import os
 
 # -----------------------------
-# Load Class Names
+# Safe model loading
 # -----------------------------
-model_path = "model/model.pth"
+possible_paths = [
+    "model/model.pth",
+    "model.pth",
+    "/app/model/model.pth",
+    "/app/model.pth"
+]
+
+model_path = None
+for p in possible_paths:
+    if os.path.exists(p):
+        model_path = p
+        break
+
+if model_path is None:
+    raise FileNotFoundError(
+        "❌ model.pth not found. Upload it to /model/model.pth or root folder."
+    )
 
 checkpoint = torch.load(model_path, map_location="cpu")
 class_names = checkpoint["class_names"]
@@ -42,7 +57,6 @@ def predict(img):
         outputs = model(img)
         probs = torch.softmax(outputs[0], dim=0)
 
-    # Return top 3 predictions
     top3_probs, top3_idxs = torch.topk(probs, 3)
     result = {class_names[i]: float(top3_probs[idx]) 
               for idx, i in enumerate(top3_idxs)}