controllers/rembg.py

from PIL import Image
from controllers.utils import image_to_bytes, bytes_to_image
from services.api_service import post_data
from controllers.fw import get_fw_query_params


def rembg_controller(image: Image, resolution: str = 'original', fw_option: str = 'none'):
    # Downscale image
    if resolution != 'original':
        if resolution == 'low':
            resolution = image.size[0] / 4
        elif resolution == 'high':
            resolution = image.size[0] / 2
        image.thumbnail((int(resolution), int(resolution)))
    # Send image
    request_data = {
        "image": image_to_bytes(image)
    }
    response = post_data(
        f"/api/rembg/{get_fw_query_params(fw_option)}", files=request_data, parse_json=False)
    if response is None:
        return None
    return bytes_to_image(response)