Update app.py

app.py

@@ -24,16 +24,10 @@ model = LitResnet()
 model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')), strict=False)
 model.eval()
 
-transform = A.Compose(
-        [
-        A.Normalize(
-                mean=(0.485, 0.456, 0.406),
-                std=(0.229, 0.224, 0.225),
-                p=1.0,
-                max_pixel_value=255,
-            ),
-        ToTensorV2()
-        ])
+transform = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
 
 inv_normalize = T.Normalize(
     mean=[-0.50/0.23, -0.50/0.23, -0.50/0.23],
@@ -50,8 +44,7 @@ def get_gradcam_image(input_tensor, label, target_layer):
 
 
 def image_classifier(input_image, top_classes=3, show_cam=True, target_layers=[2, 3], transparency=0.5):
-    input_image = transform(image=input_image)
-    #aug(image=image)
+    input_image = transform(input_image)
     input_image = input_image.unsqueeze(0)
     output = model(input_image)
     output = F.softmax(output.flatten(), dim=-1)