Update handler.py

handler.py

@@ -0,0 +1,55 @@
+from typing import Dict, List, Any
+from PIL import Image
+import io
+import base64
+from rembg import remove, new_session
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        # Preload the rembg session. 
+        # By default, it uses the 'u2net' model.
+        # This ensures the model is loaded into memory when the endpoint starts.
+        self.session = new_session()
+
+    def __call__(self, data: Dict[str, Any]) -> Any:
+        """
+        Args:
+            data (:obj: `Dict`):
+                - "inputs": PIL.Image or base64 string
+                - "parameters": (Optional) Dict containing rembg options like 'alpha_matting'
+        Return:
+            A PIL.Image or serialized image.
+        """
+        # Get inputs
+        inputs = data.get("inputs", None)
+        if inputs is None:
+            return {"error": "No inputs provided"}
+
+        # Hugging Face usually passes the image as a PIL object 
+        # if the request Content-Type is image/*
+        if isinstance(inputs, str):
+            # Handle base64 string if necessary
+            image_data = base64.b64decode(inputs)
+            image = Image.open(io.BytesIO(image_data)).convert("RGB")
+        elif isinstance(inputs, Image.Image):
+            image = inputs
+        else:
+            # Fallback for raw bytes
+            image = Image.open(io.BytesIO(inputs)).convert("RGB")
+
+        # Extract optional parameters for rembg.remove
+        # e.g., alpha_matting, bgcolor, etc.
+        params = data.get("parameters", {})
+
+        # Execute background removal
+        # rembg.remove returns a PIL Image with an alpha channel (RGBA)
+        output_image = remove(
+            image, 
+            session=self.session,
+            **params
+        )
+
+        # Return the PIL Image directly. 
+        # Hugging Face Inference Endpoints will automatically serialize 
+        # a PIL Image into a response.
+        return output_image