import os
from PIL import Image
from removebg import RemoveBg

def image_matting(old_image_path,new_image_path,api_key,colors):
    rmbg = RemoveBg(api_key,'error.log')
    rmbg.remove_background_from_img_file(old_image_path)

    parent_path = os.path.dirname(old_image_path)
    old_image_name = os.path.split(old_image_path)[-1]
    no_bg_image_name = old_image_name + '_no_bg.png'

    if no_bg_image_name in os.listdir(parent_path):
        no_bg_image_path = parent_path + "/" + no_bg_image_name
        no_bg_image = Image.open(no_bg_image_path)
        x,y = no_bg_image.size
        try:
            for color in colors:
                new_image = Image.new('RGBA',no_bg_image.size,color = color)
                new_image.paste(no_bg_image,(0,0,x,y),no_bg_image)
                path = parent_path + "/" + old_image_name + '_' + color + '.png'
                new_image.save(path)
            os.remove(no_bg_image_path)
        except Exception as e:
            print(e)
    else:
        print('error2')

def run():
    """
    API可在www.remove.bg/zh注册后免费获取
    个人用户每月可调用50次
    """
    api_key = ''
    #输入你的照片或图片路径，要具体到文件名
    old_image_path = ''
    #输出文件路径，到文件夹即可
    new_image_path = ''
    colors = ['red','blue','white']
    image_matting(old_image_path,new_image_path,api_key,colors)

if __name__ == '__main__':
    run()