import os, shutil
from PIL import Image
from removebg import RemoveBg

# 证件照背景色
Background_Color = {
    'Red': (255, 0, 0, 255),
    'Blue': (67, 142, 219, 255),
    'White': (255, 255, 255, 255)
}

pic = input("Tip：请将图片移至当前所在目录下！
请输入图片文件名(含文件后缀)：") # 图片文件名
print("当前所在目录：" + os.getcwd())

if os.path.exists('img'):
    print("
当前目录已存在img文件夹！！！")
else:
    print("
正在新建输出目录img ...")
    new_folder = '%s' % os.mkdir(os.getcwd() + '\img') # 新建文件夹
    print("输出目录img创建成功！！！
")

os.chdir('./img') # 切换至img目录下

if os.path.exists(pic) == False:
    print("
正在复制图片至img文件夹......")
    shutil.copy('../' + pic, pic)
    print("目标图片已复制到img文件夹！
")

path = '%s' % os.getcwd() # 当前所在路径
print("当前目标图片路径：	" + '%s\%s' % (path, pic)) # 输出图片所在路径

# Api每月可用50次，降低同一张图片重复抠图频率
if os.path.exists(pic + '_no_bg.png'):
    print("
当前目录已存在透明背景图片" + pic + '_no_bg.png
')
else:
    print("
正在生成错误日志(可不用理会)......")
    rmbg = RemoveBg("Your-RemoveBg-Api-Key", "error.log") # 密钥，错误日志
    print("正在进行抠图......")
    rmbg.remove_background_from_img_file('%s\%s' % (path, pic)) # 抠图
    print("已生成透明背景图片：" + pic + "_no_bg.png
")

select_bgColor = input("Tip：更换红色背景按 r，更换蓝色背景按 b，根管白色背景按 w 
请选择需要更换的证件照底色：") # 选择所需更换的底色

while True:
    if select_bgColor == 'r':
        color = 'Red'
        break
    elif select_bgColor == 'b':
        color = 'Blue'
        break
    elif select_bgColor == 'w':
        color = 'White'
        break
    else:
        select_bgColor = input("
输入有误！请重新输入：")

foreground = Image.open(pic + '_no_bg.png')  # 前景图
background = Image.new('RGBA', foreground.size, Background_Color[color])  # 背景图，大小同前景图
background.paste(foreground, mask = foreground)

print("
正在展示图片......")
background.show() # 显示图片

print("正在保存图片......
")
new_pic = pic + '_' + color + '.png' # 输出文件名
background.save(os.getcwd() + '\' + new_pic) # 保存图片

print("输出路径为：" + os.getcwd())
os.chdir('../') # 切换回原工作目录
print("
保存成功！输出文件为：" + new_pic + "

按任意键退出
") # 显示保存图片名
os.system('pause') # 按任意键继续