Update app.py

app.py

@@ -24,7 +24,19 @@ def predict(img):
     try:
         print("Starting prediction...", flush=True)
         img = PILImage.create(img)
-        pred, pred_idx, probs = learn.predict(img)
+
+        fimg = PILImage.create(img)
+        x = Resize(224)(fimg)           # -> fastai PILImage
+        x = ToTensor()(x)               # -> float tensor [0,1], shape [3,224,224]
+        x = Normalize.from_stats(*imagenet_stats)(x)
+        x = x.unsqueeze(0)              # -> [1,3,224,224]
+
+        # --- Forward pass ---
+        learn.model.eval()
+        logits = learn.model(x)
+        probs = torch.softmax(logits, dim=1)[0].cpu().numpy()
+        
+        # pred, pred_idx, probs = learn.predict(img)
         return {labels[i]: float(probs[i]) for i in range(len(labels))}
     except Exception as e:
         tb = traceback.format_exc()