Update app.py

app.py

@@ -25,16 +25,21 @@ def predict(img):
         print("Starting prediction...", flush=True)
         img = PILImage.create(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]
+        # 2) Torchvision preprocessing:
+        #    Resize shorter side to 256, center-crop 224, ToTensor, ImageNet normalize
+        tfm = T.Compose([
+            T.Resize(256),
+            T.CenterCrop(224),
+            T.ToTensor(),  # -> [C,H,W] in [0,1]
+            T.Normalize(mean=[0.485, 0.456, 0.406],
+                        std=[0.229, 0.224, 0.225]),
+        ])
+        x = tfm(img).unsqueeze(0)  # [1,3,224,224]
 
-        # --- Forward pass ---
+        # 3) Forward pass
         learn.model.eval()
-        logits = learn.model(x)
-        probs = torch.softmax(logits, dim=1)[0].cpu().numpy()
+        logits = learn.model(x)                 # [1, C]
+        probs = torch.softmax(logits, dim=1)[0] # [C]
         
         # pred, pred_idx, probs = learn.predict(img)
         return {labels[i]: float(probs[i]) for i in range(len(labels))}