Update app.py

app.py

@@ -11,6 +11,7 @@ from albumentations.pytorch.transforms import ToTensorV2
 model = models.resnet50(pretrained=False)
 model.fc = nn.Linear(2048, 21)
 model.load_state_dict(torch.load('resnet_best.pth'), strict=True)#load weights from training run
+model.eval()
 
 st.title("My awesome ML Function")
 
@@ -22,6 +23,14 @@ if uploaded_file is not None:
     st.image(image) #showing the image
     np_img = np.array(image)
     cv_image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    
+    # Let's resize the image and make a pytorch tensor from the image np array. 
+    cust_transform = A.Compose([A.Resize(height=256, width=256, p=1.0),ToTensorV2(p=1.0)], p=1.0)
+    tensor = cust_transform(image=img)
+    tensor = tensor['image'].float().resize(1,3,256,256)
+    tensor,tensor.shape
+    custom_pred = model.forward(tensor).detach().numpy() # Forward is the python method defined inside the resnet.
+    custom_pred
   elif ".csv" in uploaded_file.name: #csv check
     dataframe = pd.read_csv(uploaded_file)
     st.write(dataframe)