File size: 1,466 Bytes
44b044a 588bce7 44b044a 588bce7 44b044a 588bce7 44b044a 588bce7 44b044a b4e0815 44b044a |
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 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import torch
from PIL import Image
import io
import base64
import torchvision.transforms as T
from model import MedSAM2Model
class EndpointHandler:
def __init__(self, path=""):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = MedSAM2Model().to(self.device)
self.model.eval()
def preprocess(self, inputs):
# Unwrap if "inputs" key exists
if "inputs" in inputs:
inputs = inputs["inputs"]
image_b64 = inputs.get("image")
if not image_b64:
raise ValueError("Missing 'image' field in input.")
image_bytes = base64.b64decode(image_b64)
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
# Transform PIL image to tensor and normalize (example)
transform = T.Compose([
T.ToTensor(), # Converts to tensor and scales pixels to [0,1]
# Add normalization if your model requires it, e.g.:
# T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
image_tensor = transform(image).unsqueeze(0) # Add batch dim: [1, 3, H, W]
return image_tensor.to(self.device)
def postprocess(self, output):
return {"output": output.cpu().tolist()}
def __call__(self, inputs):
x = self.preprocess(inputs)
with torch.no_grad():
output = self.model(x)
return self.postprocess(output)
|