handler

handler.py

@@ -0,0 +1,31 @@
+import torch
+from transformers import pipeline
+from PIL import Image
+import requests
+import numpy as np
+from cv2 import imencode
+from base64 import b64encode
+
+class EndpointHandler():
+  def __init__(self):
+    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")
+    outputs = self.generator(raw_image, points_per_batch=32)
+    masks = outputs["masks"]
+    data = []
+    for index, mask in enumerate(masks):
+      cv_image = np.array(raw_image)
+      mask_image = np.zeros(cv_image.shape[:3], np.uint8)
+      mask_image[mask == False] = 255
+      retval, buffer = imencode('.png', mask_image)
+      encoded_mask = b64encode(buffer)
+      data.append({
+          "score": outputs["scores"][index].item(),
+          "mask": encoded_mask,
+          "label": ""
+      })
+    
+    return data

requirements.txt

@@ -0,0 +1,6 @@
+torch
+transformers
+pillow
+numpy
+requests
+opencv-python