Create pipeline.py

pipeline.py

@@ -0,0 +1,26 @@
+import numpy as np
+from PIL import Image
+from tensorflow.keras.models import load_model
+
+class Pipeline:
+    def __init__(self, model_path="unet_model.h5"):
+        self.model = load_model(model_path, compile=False)
+
+    def preprocess(self, image):
+        image = image.convert("L")
+        image = image.resize((176, 192))   # width, height
+        arr = np.array(image) / 255.0
+        if arr.ndim == 2:
+            arr = np.expand_dims(arr, axis=-1)
+        return np.expand_dims(arr, axis=0)
+
+    def postprocess(self, pred):
+        pred = pred[0]
+        if pred.ndim == 3 and pred.shape[-1] == 1:
+            pred = np.squeeze(pred, axis=-1)
+        return Image.fromarray((pred * 255).astype(np.uint8))
+
+    def __call__(self, image):
+        x = self.preprocess(image)
+        pred = self.model.predict(x)
+        return self.postprocess(pred)