Update pipeline.py

pipeline.py

@@ -4,8 +4,8 @@ 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 io
 import numpy as np
 from PIL import Image
 
@@ -45,8 +45,16 @@ class PreTrainedPipeline():
                     else:
                         binary_masks[f"mask_{cls}"][row][col] = 0
 
+            mask = binary_masks[f"mask_{cls}"]
+            mask *= 255
+            img = Image.fromarray(mask.astype(np.int8), mode="L")
+
+            with io.BytesIO() as out:
+                img.save(out, format="PNG")
+                png_string = out.getvalue()
+                mask = base64.b64encode(png_string).decode("utf-8")
 
-            mask_codes[f"mask_{cls}"] = base64.b64encode(binary_masks[f"mask_{cls}"])
+            mask_codes[f"mask_{cls}"] = mask
     
     
         for i in range(pred_mask.shape[-1]): #for every class
@@ -56,4 +64,4 @@ class PreTrainedPipeline():
                 "mask": mask_codes[f"mask_{i}"],
                 "score": 1.0,
             })
-            return labels
+        return labels