Update app.py

app.py

@@ -34,6 +34,19 @@ def get_model(model_name, classes, device):
         model.load_state_dict(torch.load('BaseLine-Model.pt', map_location=torch.device(device)))
 
     return model
+ 
+def get_transform(input_img, device):
+    normalize = transforms.Normalize(
+          [0.485, 0.456, 0.406], 
+          [0.229, 0.224, 0.225]
+    )
+
+    test_transform = transforms.Compose([
+        transforms.ToTensor(),
+        normalize,
+    ])
+    input_img = test_transform(input_img).unsqueeze(0).to(device)
+    return input_img
 
 def make_predictions(input_img, model_name):
     classes = ['buildings','forest', 'glacier', 'mountain', 'sea', 'street']