最近需要自制一些数据集,需要剔除背景,发现一个挺好用的工具分享给大家:
#-*-coding: utf-8 -*-
import os.path as osp
from removebg import RemoveBg
import os
from PIL import Image
# 参数填入 api key 以及批处理的日志文件位置
rmbg = RemoveBg("YOUR_API_KEY", "./error.log")
# 批处理图片的存放位置
data_path = 'IMAGES_PATH'
for pic in os.listdir(data_path):
image_path = osp.join(data_path, pic)
basename = os.path.basename(image_path).split('_')
if len(basename) > 2:
raw_path = osp.join(data_path, basename[0]+"_"+basename[1])
print("raw_path: %s" % raw_path)
del raw_path
continue
else:
remove_path = image_path + "_no_bg.png"
if osp.exists(remove_path):
print("remove_path: %s" % remove_path)
continue
print("image path: %s" % (image_path))
rmbg.remove_background_from_img_file(image_path)