Add custom handler for SAM Inference Endpoint

handler.py

@@ -0,0 +1,33 @@
+# handler.py
+import io
+import torch
+import numpy as np
+from PIL import Image
+from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
+
+# This will be called once on startup
+def initialize():
+    global mask_generator
+    # HF will mount your model files under /mnt/models
+    checkpoint = "/mnt/models/pytorch_model.bin"
+    sam = sam_model_registry["vit_b"](checkpoint=checkpoint)
+    mask_generator = SamAutomaticMaskGenerator(sam)
+
+# This handles each incoming request
+def inference(request):
+    # expect multipart/form-data with field name "image"
+    image_bytes = request.files["image"].read()
+    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    img_np = np.array(img)
+    masks = mask_generator.generate(img_np)
+
+    # combine all masks into one binary mask
+    combined = np.zeros(img_np.shape[:2], dtype=np.uint8)
+    for m in masks:
+        combined[m["segmentation"]] = 255
+
+    # serialize to PNG
+    out = io.BytesIO()
+    Image.fromarray(combined).save(out, format="PNG")
+    out.seek(0)
+    return {"mask_png": out.read()}

requirements.txt

@@ -0,0 +1,4 @@
+torch
+numpy
+Pillow
+segment-anything