Source code for odoo_xmlrpc_twisted.functions.remove_background

"""
Remove the background from an image.
"""

import os
import base64
from io import BytesIO
from PIL import Image
from rembg import remove, new_session

# Use lightweight background removal model
os.environ["U2NET_MODEL_NAME"] = "u2netp"

# ------- end of the initialisation ------- 




def remove_background(image):
    
    """
    Remove the background from an image using the 
    Python library "rembg" with the function "remove".
    See https://github.com/danielgatis/rembg.
    
    Result is the adapted image stream.
    """
      
    # Decode base64 string to bytes
    image_bytes = base64.b64decode(image)

    # Open image
    image_input = Image.open(BytesIO(image_bytes))
    
    # Remove background in a session
    try:
        # Create session with lightweight model
        session = new_session(model_name="u2netp")

        # Remove background
        image_output = remove(image_input, session=session)

    finally:
        image_input.close()  # Ensure image file is closed

    # Convert result to base64 in 3 steps:
    # 1. create a buffer to hold the image data
    buffered = BytesIO()
    # 2. save the image to the buffer in PNG format
    image_output.save(buffered, format="PNG")
    # 3. convert the image data in the buffer to a base64-encoded string     
    image_result = base64.b64encode(buffered.getvalue())
    
    return(image_result)