infer/rembg.py

from rembg import remove, new_session
from .utils import timing_decorator

class Removebg():
    def __init__(self, name="u2net"):
        '''
            name: rembg
        '''
        self.session = new_session(name)

    @timing_decorator("remove background")
    def __call__(self, rgb_img, force=False):
        '''
            inputs:
                rgb_img: PIL.Image, with RGB mode expected
                force: bool, input is RGBA mode
            return:
                rgba_img: PIL.Image with RGBA mode
        '''
        if rgb_img.mode == "RGBA":
            if force:
                rgb_img = rgb_img.convert("RGB")
            else:
                return rgb_img
        rgba_img = remove(rgb_img, session=self.session)
        return rgba_img