rembg / handler.py

Update handler.py

2e68455 verified 9 months ago

2.04 kB

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.
        # model_name = "u2net"
        model_name = "u2netp"
        # model_name = "birefnet-general-lite"
        self.session = new_session(model_name)

    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

	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.
	# model_name = "u2net"
	model_name = "u2netp"
	# model_name = "birefnet-general-lite"
	self.session = new_session(model_name)

	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