commit 57cd9a1f392b63a1de61edb225cfabbb6eac1424 Author: douboer Date: Tue Jan 20 13:44:28 2026 +0800 first commit diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000..a9c3e12 Binary files /dev/null and b/.DS_Store differ diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..2de9ca2 --- /dev/null +++ b/.gitignore @@ -0,0 +1,36 @@ +# Build and Release Folders +bin-debug/ +bin-release/ +[Oo]bj/ +[Bb]in/ + +# Other files and folders +.settings/ + +node_modules/ +public/ +myenv/ +venv/ +package.json +package-lock.json + +# Executables +*.swf +*.air +*.ipa +*.apk + +# Project files, i.e. `.project`, `.actionScriptProperties` and `.flexProperties` +# should NOT be excluded as they contain compiler settings and other important +# information for Eclipse / Flash Builder. + +*.mov +*.mp4 +cert/nginx/*.crt +cert/nginx/*.key + +*.jpg +*.heic +*.jpeg +*.png + diff --git a/README.md b/README.md new file mode 100644 index 0000000..105efc3 --- /dev/null +++ b/README.md @@ -0,0 +1,208 @@ +# 图片去背景工具 + +使用rembg库实现的Python去背景工具 + +## 快速开始 + +```bash +# 激活虚拟环境 +source ~/venv/bin/activate + +# 使用默认参数处理images文件夹 +python remove_background.py + +# 处理单个文件 +python remove_background.py input.jpg output.png + +# 查看所有参数-m, --model) + +不同模型适用于不同场景： + +| 模型名称 | 大小 | 适用场景 | 推荐度 | +|---------|------|---------|--------| +| **isnet-general-use** | 179MB | 通用场景 | ⭐⭐⭐⭐⭐ 默认推荐 | +| birefnet-general | 250MB | 通用场景，质量更高 | ⭐⭐⭐⭐⭐ | +| birefnet-portrait | 250MB | 人像专用 | ⭐⭐⭐⭐⭐ | +| u2net | 176MB | 经典通用模型 | ⭐⭐⭐⭐ | +| u2netp | 4.7MB | 快速处理 | ⭐⭐⭐ | +| u2net_human_seg | 176MB | 人物分割 | ⭐⭐⭐⭐ | +| isnet-anime | 179MB | 动漫角色 | ⭐⭐⭐⭐ | +| silueta | 43MB | 精简快速 | ⭐⭐⭐ | + +**使用示例**: +```bash +# 使用默认模型 +python remove_background.py input.jpg + +# 使用人像专用模型 +python remove_background.py input.jpg output.png -m birefnet-portrait + +# 使用快速模型 +python remove_background.py input.jpg output.png -m u2netp +`` +```bash +ls -lh ~/.u2net/ +``` + +### 模型大小参考 + +| 模型名称 | 文件大小 | 特点 | +|---------|---------|------| +| u2net | 176MB | 通用模型 | +| u2netp | 4.7MB | 轻量级，速度快 | +| isnet-general-use | 179MB | 新一代通用，推荐 | +| birefnet-general | ~250MB | 最新通用模型 | +| birefnet-portrait | ~250MB | 人像专用 | + +### 手动下载模型（网络问题时） + +```bash +# 创建目录 +mkdir -p ~/.u2net/ + +# 下载指定模型（以isnet-general-use为例） +curl -L "https://github.com/danielgatis/rembg/releases/download/v0.0.0/isnet-general-use.onnx" \ + -o ~/.u2net/isnet-general-use.onnx +``` + +### 清理模型缓存 + +```bash +# 删除所有已下载的模型 +rm -rf ~/.u2net/ + +# 删除特定模型 +rm ~/.u2net/u2net.onnx +``` + +## 可调整参数说明 + +### 1. 模型选择 (model_name) + +不同模型适用于不同场景： + +- **u2net** (默认): 通用模型，适合大多数场景 +- **u2netp**: 轻量版，速度更快但精度稍低 +- **u2net_human_seg**: 专门用于人物分割 +- **silueta**: 精简版u2net (43MB)，速度快 +- **isnet-general-use**: 新一代通用模型，效果可能更好 +- **isnet-anime**: 专门用于动漫角色 +- **birefnet-general**: 最新的通用模型，推荐尝试 +- **birefnet-portrait**: 专门用于人像 +- **birefnet-general-lite**: 轻量版birefnet + +**建议**: 如果u2net效果不好，试试 `isnet-general-use` 或 `birefnet-general` + +### 2. Alpha Matting 参数 + +Alpha Matting 是后处理步骤，可以显著改善边缘质量，特别是头发、毛发等细节。 + +#### alpha_matting (True/False) +- **作用**: 是否启用alpha matting +- **默认**: False +- **建议**: 如果边缘不自然，启用此选项 + +#### alpha_matting_foreground_threshold (0-255) +- **作用**: 前景阈值，控制哪些区域被认为是前景 +- **默认**: 240 +- **调整建议**: + - 值越大(如270): 保留更多细节，但可能保留一些背景 + - 值越小(如210): 去除更彻底，但可能丢失细节 + - 如果前景被过度去除，增加此值 + - 如果背景残留太多，减小此值 + +#### alpha_matting_background_threshold (0-255) +- **作用**: 背景阈值，控制哪些区域被认为是背景 +- **默认**: 10 +- **调整建议**: + - 值越大(如20-30): 去除背景更彻底 + - 值越小(如5): 保留更多过渡区域 + - 如果背景残留，增加此值 + +#### alpha_matting_erode_size (像素) +- **作用**: 侵蚀大小，用于平滑边缘 +- **默认**: 10 +- **调整建议**: + - 值越大(如15-20): 边缘更平滑，但可能损失细节 + - 值越小(如5-8): 保留更多细节，但边缘可能不够平滑 + +### 3. Mask后处理 (post_process_mask) + +- **作用**: 对mask进行额外的后处理 +- **默认**: False +- **建议**: 可以尝试启用看效果是否改善 + +## 常见问题解决 + +### 问题1: 前景被过度去除 +**解决方案**: +```python +alpha_matting = True +alpha_matting_foreground_threshold = 270 # 增加此值 +alpha_matting_background_threshold = 10 # 保持较小 +``` + +### 问题2: 背景残留太多 +**解决方案**: +```python +alpha_matting = True +alpha_matting_foreground_threshold = 240 # 保持默认或减小 +alpha_matting_background_threshold = 20 # 增加此值 +post_process_mask = True # 启用后处理 +``` + +### 问题3: 边缘不自然、有锯齿 +**解决方案**: +```python +alpha_matting = True +alpha_matting_erode_size = 15 # 增加平滑程度 +``` + +### 问题4: 毛发、头发细节丢失 +**解决方案**: +```python +model_name = "birefnet-portrait" # 使用人像专用模型 +alpha_matting = True +alpha_matting_foreground_threshold = 270 # 增加以保留细节 +alpha_matting_erode_size = 5 # 减小以保留细节 +``` + +## 推荐配置 + +### 配置1: 高质量人像 +```python +model_name = "birefnet-portrait" +alpha_matting = True +alpha_matting_foreground_threshold = 260 +alpha_matting_background_threshold = 15 +alpha_matting_erode_size = 10 +post_process_mask = True +``` + +### 配置2: 通用高质量 +```python +model_name = "birefnet-general" +alpha_matting = True +alpha_matting_foreground_threshold = 250 +alpha_matting_background_threshold = 12 +alpha_matting_erode_size = 10 +post_process_mask = True +``` + +### 配置3: 快速处理 +```python +model_name = "u2netp" +alpha_matting = False +post_process_mask = False +``` + +## 测试不同参数 + +建议按以下顺序调整: + +1. 先尝试不同的模型 +2. 启用alpha_matting +3. 调整foreground_threshold和background_threshold +4. 最后调整erode_size + +每次修改后运行脚本，对比结果。 diff --git a/output/.DS_Store b/output/.DS_Store new file mode 100644 index 0000000..5008ddf Binary files /dev/null and b/output/.DS_Store differ diff --git a/remove_background.py b/remove_background.py new file mode 100644 index 0000000..43531c2 --- /dev/null +++ b/remove_background.py @@ -0,0 +1,258 @@ +""" +图片去背景工具 +使用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) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..6844ac3 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +rembg[gpu] +pillow +pillow-heif