Add base64

handler.py

@@ -4,7 +4,8 @@ from PIL import Image
 from huggingface_hub import hf_hub_download
 import torchvision.transforms as T
 from typing import Dict, List, Any
-
+from io import BytesIO
+import base64
 # image = Image.open("man.png").convert("RGB")
 # trimap = Image.open("mask2.png").convert("L")
 
@@ -17,8 +18,13 @@ class EndpointHandler():
             "hustvl/vitmatte-small-composition-1k")
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        image = data.pop("image")
-        trimap = data.pop("trimap")
+        inputs = data.pop("inputs", data)
+        # parameters = data.pop("parameters", {"mode": "image"})
+
+        image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
+        trimap = Image.open(BytesIO(base64.b64decode(inputs['trimap'])))
+        # image = data.pop("image")
+        # trimap = data.pop("trimap")
 
         inputs = self.processor(
             images=image, trimaps=trimap, return_tensors="pt")