Create new file

make_predictions.py

@@ -0,0 +1,20 @@
+import torch
+import numpy as np
+import cv2
+
+def make_predictions(model, input_img, threshold = 0.05):
+    model.eval()
+    with torch.no_grad():
+        image = input_img
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        image = image.astype("float32") / 255.0
+        image = cv2.resize(image, (height, width))
+        image = np.transpose(image, (2, 0, 1))
+        image = np.expand_dims(image, 0)
+        image = torch.from_numpy(image).to(device)
+        predMask = model(image).squeeze()
+        predMask = torch.sigmoid(predMask)
+        predMask = predMask.cpu().numpy()
+        predMask = (predMask > threshold) * 255
+        predMask = predMask.astype(np.uint8)
+        return predMask