""" 图片去背景工具 使用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)