Fix image upload errors

handler.py

@@ -29,10 +29,11 @@ class EndpointHandler():
             self.model = SamModel.from_pretrained("facebook/sam-vit-base").to(device)
             self.processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def __call__(self, data):
         """
         Called on every HTTP request.
         Handles both base64-encoded images and PIL images.
+        Returns a PIL Image object.
         """
         # 1. Parse and decode the input image
         inputs = data.pop("inputs", None)
@@ -41,12 +42,10 @@ class EndpointHandler():
 
         # Check the type of inputs to handle both base64 strings and pre-processed PIL Images
         if isinstance(inputs, Image.Image):
-            # Input is already a PIL Image
             img = inputs.convert("RGB")
         elif isinstance(inputs, str):
-            # Input is a base64-encoded string
             if inputs.startswith("data:"):
-                inputs = inputs.split(",", 1)[1]  # Handle data URL format
+                inputs = inputs.split(",", 1)[1]
             image_bytes = base64.b64decode(inputs)
             img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
         else:
@@ -91,36 +90,27 @@ class EndpointHandler():
             size = min(width, height) // 8
             mask_binary[center_y-size:center_y+size, center_x-size:center_x+size] = 255
         
-        # 5. Encode the output image to base64
-        out = io.BytesIO()
-        Image.fromarray(mask_binary).save(out, format="PNG")
-        out.seek(0)
-        mask_base64 = base64.b64encode(out.getvalue()).decode('utf-8')
-        
-        # 6. Return the response
-        return [{"mask_png_base64": mask_base64}]
+        # 5. Create and return the output PIL Image
+        output_img = Image.fromarray(mask_binary)
+        return output_img
 
 def main():
-    # This main function shows how a client would call the endpoint.
+    # This main function shows how a client would call the endpoint locally.
     input_path = "/Users/rp7/Downloads/test.jpeg"
     output_path = "output.png"
 
-    # 1. Prepare the payload
+    # 1. Prepare the payload with a base64-encoded image string
     with open(input_path, "rb") as f:
         img_bytes = f.read()
     img_b64 = base64.b64encode(img_bytes).decode("utf-8")
     payload = {"inputs": "data:image/jpeg;base64,{}".format(img_b64)}
 
-    # 2. Instantiate handler and call it
+    # 2. Instantiate handler and get the PIL Image result
     handler = EndpointHandler(path=".")
-    result = handler(payload)
+    result_img = handler(payload)
 
-    # 3. Process the response
-    mask_b64 = result[0]["mask_png_base64"]
-    mask_bytes = base64.b64decode(mask_b64)
-    
-    mask_img = Image.open(io.BytesIO(mask_bytes))
-    mask_img.save(output_path)
+    # 3. Save the resulting image
+    result_img.save(output_path)
     print("Wrote mask to {}".format(output_path))
 
 if __name__ == "__main__":