Update app.py

app.py

@@ -58,6 +58,14 @@ classes = ['Acanthostichus',
            'Zasphinctus']
 class_to_idx = {idx: cls for idx, cls in enumerate(classes)}
 
+train_transforms = torchvision.transforms.Compose(
+    [
+        torchvision.transforms.ToTensor(),
+        torchvision.transforms.Resize((TARGET_SIZE[0], TARGET_SIZE[1]), antialias=True),
+        CMYKToRGB(),
+    ]
+)
+
 tf = Compose([torchvision.transforms.Resize((64, 64), antialias=True)])
 
 model = ViTIJEPA(64, 4, 3, 64, 8, 8, len(classes))
@@ -65,16 +73,15 @@ model.load_state_dict(torch.load("vit_ijepa_ant_1.pt", map_location=torch.device
 
 
 def ant_genus_classification(image):
-    image = torch.Tensor(image)
-    image = image.unsqueeze(0)
-    image = rearrange(image, 'b h w c -> b c h w')
-    image = tf(image)
+    image = train_transforms(image).unsqueeze(0)
 
     print(image.shape)
     with torch.no_grad():
-        prediction = torch.nn.functional.softmax(model(image)[0], dim=0)
-        # print(prediction.tolist())
-        confidences = {class_to_idx[i]: float(prediction[i]) for i in range(len(classes))}
+        y_hat = model(x)
+        y_hat = torch.nn.functional.softmax(y_hat, dim=1)
+        preds = torch.argmax(y_hat, dim=1)
+        
+        confidences = {class_to_idx[i]: float(preds[i]) for i in range(len(classes))}
     return confidences
     # prediction = model(image)[0]
     # prediction = prediction.tolist()