Add custom handler for SAM2

handler.py

@@ -0,0 +1,41 @@
+from typing import Dict, List, Any
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+import torch
+import numpy as np
+from PIL import Image
+import io
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        # Initialize SAM2 predictor with small model
+        self.predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-small")
+        
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Args:
+            data: Dictionary with "inputs" key containing image bytes
+        Returns:
+            Dictionary containing masks and scores
+        """
+        # Get input image
+        if "inputs" not in data:
+            raise ValueError("No inputs provided")
+            
+        # Convert input image bytes to PIL Image
+        image = Image.open(io.BytesIO(data["inputs"]))
+        image = np.array(image)
+        
+        # Process with SAM2
+        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+            self.predictor.set_image(image)
+            masks, scores, _ = self.predictor.predict()
+            
+        # Convert masks to lists for JSON serialization
+        if masks is not None:
+            masks = [mask.tolist() for mask in masks]
+            scores = scores.tolist() if scores is not None else None
+            
+        return {
+            "masks": masks,
+            "scores": scores
+        }

requirements.txt

@@ -0,0 +1 @@
+sam2