change output format

handler.py

@@ -1,10 +1,26 @@
 import base64
+import gzip
+import numpy as np
 from io import BytesIO
 from typing import Dict, List, Any
 from PIL import Image
 import torch
 from transformers import SamModel, SamProcessor
 
+def pack_bits(boolean_tensor):
+    # Flatten the tensor and add padding if necessary
+    flat = boolean_tensor.flatten()
+    if flat.size()[0] % 8 != 0:
+        padding = np.zeros((8 - flat.size % 8,), dtype=bool)
+        flat = np.concatenate([flat, padding])
+    
+    # Reshape into bytes and pack into binary string
+    packed = np.packbits(flat.reshape((-1, 8)))
+    packed = packed.tobytes()
+    return gzip.compress(packed)
+# json_str = json.dumps({"shape": boolean_tensor.shape, "data": binary_str})
+
+
 class EndpointHandler():
     def __init__(self, path=""):
         # Preload all the elements you are going to need at inference.
@@ -30,7 +46,13 @@ class EndpointHandler():
 
         model_inputs = self.processor(image, input_points=input_points, return_tensors="pt").to(self.device)
         outputs = self.model(**model_inputs)
-        masks = self.processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), model_inputs["original_sizes"].cpu(), model_inputs["reshaped_input_sizes"].cpu())
+        masks = self.processor.image_processor.post_process_masks(
+            outputs.pred_masks.cpu(), 
+            model_inputs["original_sizes"].cpu(), 
+            model_inputs["reshaped_input_sizes"].cpu())
         scores = outputs.iou_scores
 
-        return {"masks": masks, "scores": scores}
+        packed = [base64.b64encode(pack_bits(masks[0][0][i])).decode() for i in range(masks[0].shape[1])]
+        shape = list(masks[0].shape)[2:]
+
+        return {"masks": packed, "scores": scores[0][0].tolist(), "shape": shape}