File size: 1,079 Bytes
4a0358b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 | from typing import Dict, Any
from transformers import pipeline
import torch
from PIL import Image
import numpy as np
class EndpointHandler:
def __init__(self, path: str):
"""
Initialize the handler, load the SAM2 model.
"""
# Load SAM2 model and prepare pipeline
self.model = pipeline("image-segmentation", model=path)
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
"""
Perform inference on the input data and return the result.
"""
image_data = data["inputs"] # Assuming inputs key contains the image
point_coords = data["point_coords"]
point_labels = data["point_labels"]
# Convert image from base64 or other formats to numpy array if necessary
# Assuming image_data is already in a suitable format for model inference
# Running the inference with SAM2 model
segmentation_result = self.model(image_data)
return {"result": segmentation_result} # Return results in a dictionary
|