pipelines/remove_background.py

import io
from typing import Union

from PIL import Image
from rembg import remove
from util.commons import read_url


class RemoveBackground:
    def remove(self, image: Union[str, Image.Image]) -> Image.Image:
        if type(image) is str:
            image = Image.open(io.BytesIO(read_url(image)))

        output = remove(image)
        return output