Upload 2 files

handler.py

@@ -2,6 +2,7 @@ import torch
 from PIL import Image
 import io
 import base64
+import torchvision.transforms as T
 
 from model import MedSAM2Model
 
@@ -11,7 +12,12 @@ class EndpointHandler:
         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.")
@@ -19,10 +25,14 @@ class EndpointHandler:
         image_bytes = base64.b64decode(image_b64)
         image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
 
-        # Convert PIL Image to PyTorch tensor (C x H x W), normalized [0,1]
-        image_tensor = torch.from_numpy(
-            (np.array(image) / 255.0).astype('float32')
-        ).permute(2, 0, 1).unsqueeze(0)  # Add batch dimension
+        # 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)
 

requirements.txt

@@ -1,3 +1,4 @@
 torch
+torchvision
 Pillow
 huggingface_hub