Update pipeline.py

pipeline.py

@@ -3,67 +3,57 @@ from typing import Any, Dict, List
 
 import tensorflow as tf
 from tensorflow import keras
+from huggingface_hub import from_pretrained_keras, hf_hub_download
 from PIL import Image
 import base64
+import numpy as np
+from PIL import Image
 
 
-MODEL_FILENAME = "saved_model.pb"
-CONFIG_FILENAME = "config.json"
-
 
-class PreTrainedPipeline(Pipeline):
+class PreTrainedPipeline():
     def __init__(self, model_id: str):
 
+        self.model = keras.models.load_model(model_id)
+
+    def __call__(self, inputs: "Image.Image")-> List[Dict[str, Any]]:
+
+        with Image.open(inputs) as img:
+            img = np.array(img)
+
+        im = tf.image.resize(img, (128, 128))
+        im = tf.cast(im, tf.float32) / 255.0
+        pred_mask = model.predict(im[tf.newaxis, ...])
+        pred_mask_arg = tf.argmax(pred_mask, axis=-1)
 
-        # Reload Keras SavedModel
-        self.model = keras.models.load_model('./model.h5')
-
-        # Number of labels
-        self.num_labels = self.model.output_shape[1]
-
-        self.id2label = self.id2label = {"0": "pet", "1":"no_pet"}
-
-    def __call__(self, inputs: "Image.Image") -> List[Dict[str, Any]]:
-        """
-        Args:
-            inputs (:obj:`PIL.Image`):
-                The raw image representation as PIL.
-                No transformation made whatsoever from the input. Make all necessary transformations here.
-        Return:
-            A :obj:`list`:. The list contains items that are dicts should be liked {"label": "XXX" (str), mask: "base64 encoding of the mask" (str), "score": float}
-                It is preferred if the returned list is in decreasing `score` order
-        """
-        # Resize image to expected size
-
-        expected_input_size = self.model.input_shape
-
-        
-        with Image.open(inputs) as im:
-            inputs = np.array(im)
-        if expected_input_size[-1] == 1:  
-            inputs = inputs.convert("L")
-
-
-        
-        target_size = (expected_input_size[1], expected_input_size[2])
-        img = tf.image.resize(inputs, target_size)
-        
-        img_array = tf.keras.preprocessing.image.img_to_array(img)
-        img_array = img_array[tf.newaxis, ...]
-        
-        predictions = self.model.predict(img_array)
-
-        self.single_output_unit = (
-            self.model.output_shape[1] == 1
-        )  # if there are two classes
-        
         labels = []
-        for i in enumerate(predictions):
+
+        binary_masks = {}
+        mask_codes = {}
+
+
+        for cls in range(pred_mask.shape[-1]):
+
+            binary_masks[f"mask_{cls}"] = np.zeros(shape = (pred_mask.shape[1], pred_mask.shape[2]))
+            for row in range(pred_mask_arg[0][1].get_shape().as_list()[0]):
+
+                for col in range(pred_mask_arg[0][2].get_shape().as_list()[0]):
+
+                    if pred_mask_arg[0][row][col] == cls:
+                        
+                        binary_masks[f"mask_{cls}"][row][col] = 1
+                    else:
+                        binary_masks[f"mask_{cls}"][row][col] = 0
+
+
+            mask_codes[f"mask_{cls}"] = base64.b64encode(binary_masks[f"mask_{cls}"])
+    
+    
+        for i in range(pred_mask.shape[-1]): #for every class
             
             labels.append({
-                "label": str(i[0]),
-                "mask": base64.b64encode(i[1]),
+                "label": f"LABEL_{i}",
+                "mask": mask_codes[f"mask_{i}"],
                 "score": 1.0,
             })
-
-        return sorted(labels)
+            return labels