259 lines
9.9 KiB
Python
259 lines
9.9 KiB
Python
"""
|
||
图片去背景工具
|
||
使用rembg库自动去除图片背景
|
||
"""
|
||
import os
|
||
import argparse
|
||
from pathlib import Path
|
||
from rembg import remove, new_session
|
||
from PIL import Image
|
||
|
||
# 支持HEIC格式
|
||
try:
|
||
from pillow_heif import register_heif_opener
|
||
register_heif_opener()
|
||
HEIC_SUPPORTED = True
|
||
except ImportError:
|
||
HEIC_SUPPORTED = False
|
||
|
||
def str2bool(v):
|
||
"""将字符串转换为布尔值"""
|
||
if isinstance(v, bool):
|
||
return v
|
||
if v.lower() in ('yes', 'true', 't', 'y', '1'):
|
||
return True
|
||
elif v.lower() in ('no', 'false', 'f', 'n', '0'):
|
||
return False
|
||
else:
|
||
raise argparse.ArgumentTypeError('布尔值应为true或false')
|
||
|
||
def remove_background(input_path, output_path, session=None, **kwargs):
|
||
"""
|
||
去除图片背景
|
||
|
||
Args:
|
||
input_path: 输入图片路径
|
||
output_path: 输出图片路径
|
||
session: rembg会话对象（可选）
|
||
**kwargs: 其他参数，如alpha_matting相关参数
|
||
"""
|
||
print(f"正在处理: {input_path}")
|
||
|
||
# 读取输入图片
|
||
input_image = Image.open(input_path)
|
||
|
||
# 去除背景
|
||
output_image = remove(input_image, session=session, **kwargs)
|
||
|
||
# 保存输出图片
|
||
output_image.save(output_path)
|
||
|
||
print(f"已保存: {output_path}")
|
||
|
||
def process_images_folder(input_folder, output_folder, model_name="u2net",
|
||
alpha_matting=False, alpha_matting_foreground_threshold=240,
|
||
alpha_matting_background_threshold=10, alpha_matting_erode_size=10,
|
||
post_process_mask=False):
|
||
"""
|
||
批量处理文件夹中的所有图片
|
||
|
||
Args:
|
||
input_folder: 输入文件夹路径
|
||
output_folder: 输出文件夹路径
|
||
model_name: 模型名称，可选值:
|
||
- u2net (默认): 通用模型
|
||
- u2netp: 轻量版u2net
|
||
- u2net_human_seg: 人物分割
|
||
- silueta: 精简版u2net (43MB)
|
||
- isnet-general-use: 新的通用模型
|
||
- isnet-anime: 动漫角色高精度分割
|
||
- birefnet-general: 通用模型
|
||
- birefnet-portrait: 人像模型
|
||
alpha_matting: 是否启用alpha matting后处理（改善边缘质量）
|
||
alpha_matting_foreground_threshold: 前景阈值 (0-255)，值越大保留越多前景
|
||
alpha_matting_background_threshold: 背景阈值 (0-255)，值越大去除越多背景
|
||
alpha_matting_erode_size: 侵蚀大小，用于平滑边缘
|
||
post_process_mask: 是否启用mask后处理
|
||
"""
|
||
# 创建输出文件夹
|
||
Path(output_folder).mkdir(parents=True, exist_ok=True)
|
||
|
||
# 创建会话（重用会话可以提高性能）
|
||
print(f"使用模型: {model_name}")
|
||
session = new_session(model_name)
|
||
|
||
# 支持的图片格式
|
||
image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.webp'}
|
||
if HEIC_SUPPORTED:
|
||
image_extensions.update({'.heic', '.heif'})
|
||
else:
|
||
print("提示: 未安装pillow-heif，HEIC格式不可用。安装方法: pip install pillow-heif")
|
||
|
||
# 获取所有图片文件
|
||
input_path = Path(input_folder)
|
||
image_files = [f for f in input_path.iterdir()
|
||
if f.is_file() and f.suffix.lower() in image_extensions]
|
||
|
||
if not image_files:
|
||
print(f"在 {input_folder} 中没有找到图片文件")
|
||
return
|
||
|
||
print(f"找到 {len(image_files)} 张图片，开始处理...")
|
||
print(f"Alpha Matting: {'启用' if alpha_matting else '禁用'}")
|
||
if alpha_matting:
|
||
print(f" - 前景阈值: {alpha_matting_foreground_threshold}")
|
||
print(f" - 背景阈值: {alpha_matting_background_threshold}")
|
||
print(f" - 侵蚀大小: {alpha_matting_erode_size}")
|
||
print(f"Mask后处理: {'启用' if post_process_mask else '禁用'}")
|
||
print("-" * 50)
|
||
|
||
# 处理每张图片
|
||
for i, image_file in enumerate(image_files, 1):
|
||
try:
|
||
# 输出文件名（保持原始名称，改为PNG格式以支持透明背景）
|
||
output_filename = image_file.stem + '_nobg.png'
|
||
output_path = Path(output_folder) / output_filename
|
||
|
||
print(f"[{i}/{len(image_files)}] ", end="")
|
||
remove_background(
|
||
str(image_file),
|
||
str(output_path),
|
||
session=session,
|
||
alpha_matting=alpha_matting,
|
||
alpha_matting_foreground_threshold=alpha_matting_foreground_threshold,
|
||
alpha_matting_background_threshold=alpha_matting_background_threshold,
|
||
alpha_matting_erode_size=alpha_matting_erode_size,
|
||
post_process_mask=post_process_mask
|
||
)
|
||
|
||
except Exception as e:
|
||
print(f"处理 {image_file.name} 时出错: {e}")
|
||
|
||
print("-" * 50)
|
||
print(f"处理完成！结果保存在 {output_folder} 文件夹中")
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(
|
||
description='图片去背景工具 - 使用rembg自动去除图片背景',
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog='''
|
||
示例用法:
|
||
# 使用默认参数处理images文件夹
|
||
python remove_background.py
|
||
|
||
# 处理单个文件
|
||
python remove_background.py input.jpg output.png
|
||
|
||
# 处理指定文件夹
|
||
python remove_background.py my_images/ my_output/
|
||
|
||
# 使用不同模型
|
||
python remove_background.py input.jpg output.png -m birefnet-portrait
|
||
|
||
# 自定义alpha matting参数
|
||
python remove_background.py input.jpg output.png -ft 260 -bt 12 -es 5
|
||
'''
|
||
)
|
||
|
||
# 必需参数
|
||
parser.add_argument('input', nargs='?', default='images',
|
||
help='输入文件或文件夹路径（默认: images）')
|
||
parser.add_argument('output', nargs='?', default=None,
|
||
help='输出文件或文件夹路径（可选，默认为output/）')
|
||
|
||
# 模型选择
|
||
parser.add_argument('-m', '--model', default='isnet-general-use',
|
||
choices=['u2net', 'u2netp', 'u2net_human_seg', 'silueta',
|
||
'isnet-general-use', 'isnet-anime',
|
||
'birefnet-general', 'birefnet-general-lite',
|
||
'birefnet-portrait', 'birefnet-dis',
|
||
'birefnet-hrsod', 'birefnet-cod', 'birefnet-massive'],
|
||
help='选择使用的模型 (默认: isnet-general-use)')
|
||
|
||
# Alpha Matting参数
|
||
parser.add_argument('-a', '--alpha-matting', type=str2bool, nargs='?', const=True, default=True,
|
||
metavar='true/false',
|
||
help='启用alpha matting后处理（默认: true）。用法: -a 或 -a true 或 -a false')
|
||
parser.add_argument('-ft', '--foreground-threshold', type=int, default=245,
|
||
help='前景阈值 (0-255)，值越大保留越多细节 (默认: 245)')
|
||
parser.add_argument('-bt', '--background-threshold', type=int, default=8,
|
||
help='背景阈值 (0-255)，值越大去除越多背景 (默认: 8)')
|
||
parser.add_argument('-es', '--erode-size', type=int, default=2,
|
||
help='侵蚀大小，用于平滑边缘，值越大越平滑但可能丢失细节 (默认: 2)')
|
||
|
||
# 其他选项
|
||
parser.add_argument('-p', '--post-process', type=str2bool, nargs='?', const=True, default=True,
|
||
metavar='true/false',
|
||
help='启用mask后处理（默认: true）。用法: -p 或 -p true 或 -p false')
|
||
|
||
args = parser.parse_args()
|
||
|
||
print("=" * 50)
|
||
print("图片去背景工具")
|
||
print("=" * 50)
|
||
|
||
# 判断输入是文件还是文件夹
|
||
input_path = Path(args.input)
|
||
|
||
if not input_path.exists():
|
||
print(f"错误: 输入路径不存在: {args.input}")
|
||
exit(1)
|
||
|
||
# 处理单个文件
|
||
if input_path.is_file():
|
||
# 确定输出路径
|
||
if args.output is None:
|
||
output_path = input_path.parent / 'output' / (input_path.stem + '_nobg.png')
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
else:
|
||
output_path = Path(args.output)
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
print(f"输入文件: {input_path}")
|
||
print(f"输出文件: {output_path}")
|
||
print(f"模型: {args.model}")
|
||
print(f"Alpha Matting: {'启用' if args.alpha_matting else '禁用'}")
|
||
if args.alpha_matting:
|
||
print(f" - 前景阈值: {args.foreground_threshold}")
|
||
print(f" - 背景阈值: {args.background_threshold}")
|
||
print(f" - 侵蚀大小: {args.erode_size}")
|
||
print(f"Mask后处理: {'启用' if args.post_process else '禁用'}")
|
||
print("-" * 50)
|
||
|
||
# 创建会话
|
||
session = new_session(args.model)
|
||
|
||
# 处理图片
|
||
remove_background(
|
||
str(input_path),
|
||
str(output_path),
|
||
session=session,
|
||
alpha_matting=args.alpha_matting,
|
||
alpha_matting_foreground_threshold=args.foreground_threshold,
|
||
alpha_matting_background_threshold=args.background_threshold,
|
||
alpha_matting_erode_size=args.erode_size,
|
||
post_process_mask=args.post_process
|
||
)
|
||
|
||
print("-" * 50)
|
||
print(f"处理完成！结果保存在: {output_path}")
|
||
|
||
# 处理文件夹
|
||
elif input_path.is_dir():
|
||
output_folder = args.output if args.output else 'output'
|
||
|
||
process_images_folder(
|
||
str(input_path),
|
||
output_folder,
|
||
model_name=args.model,
|
||
alpha_matting=args.alpha_matting,
|
||
alpha_matting_foreground_threshold=args.foreground_threshold,
|
||
alpha_matting_background_threshold=args.background_threshold,
|
||
alpha_matting_erode_size=args.erode_size,
|
||
post_process_mask=args.post_process
|
||
)
|
||
|
||
else:
|
||
print(f"错误: 不支持的输入类型: {args.input}")
|
||
exit(1)
|