handler

handler.py

@@ -11,8 +11,9 @@ class EndpointHandler():
     self.device = "cuda" if torch.cuda.is_available() else "cpu"
     self.generator = pipeline("mask-generation", model="facebook/sam-vit-large", device=self.device)
 
-  def __call__(self, url):
-    raw_image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
+  def __call__(self, data):
+    inputs = data.pop("inputs", data)
+    raw_image = Image.open(requests.get(inputs, stream=True).raw).convert("RGB")
     outputs = self.generator(raw_image, points_per_batch=32)
     masks = outputs["masks"]
     data = []