gavin/remove_backgroud
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update at 2026-03-28 19:36:19

Browse Source
This commit is contained in:
douboer
2026-03-28 19:36:19 +08:00
parent bd0e4a660a
commit 278f1e0f69
6 changed files with 1039 additions and 389 deletions
Download Patch File Download Diff File Expand all files Collapse all files

BIN
.DS_Store vendored
View File

Binary file not shown.

497
README.md
View File

@@ -1,293 +1,294 @@
# 图片去背景工具
使用rembg库实现的Python去背景工具
用于书画、篆刻作品的去背景与去主体补背景脚本。
## 快速开始
默认模式优先走书画专用前景提取逻辑：
- 对墨色笔画做局部背景校正与明暗差分
- 对印章做更严格的红色检测
- 在粗掩码内进一步细化 alpha，尽量减少字边残留底色
- 仅在需要时才回退到 `rembg`
## 功能概览
- 单张图片去背景，输出透明 PNG
- 批量处理单层目录中的图片
- 支持 `artwork`、`auto`、`rembg` 三种前景提取模式
- 支持 `calligraphy`、`seal`、`auto` 三种书画类型
- 支持 AOT-GAN 去主体补背景
- 支持常见格式：`jpg`、`jpeg`、`png`、`bmp`、`webp`
- 安装 `pillow-heif` 后可读取 `heic` / `heif`
## 依赖
项目当前依赖见 [requirements.txt](/Users/gavin/removeback/requirements.txt)：
```txt
rembg[gpu]
pillow
pillow-heif
opencv-python
```
如果要使用 AOT-GAN 补背景，还需要额外安装 `torch` / `torchvision`，并准备 AOT-GAN 代码目录与权重文件。
## 环境准备
推荐使用现有虚拟环境：
```bash
# 激活虚拟环境
source ~/venv/bin/activate
# 使用默认参数处理images文件夹
python remove_background.py
# 处理单个文件
python remove_background.py input.jpg output.png
# 查看所有参数
python remove_background.py -h
```
不同模型适用于不同场景：
安装依赖：
| 模型名称 | 大小 | 适用场景 | 推荐度 |
|---------|------|---------|--------|
| **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
# 书画类文字偏浅的温和参数示例
python remove_background.py images output \
--remove-subject --black_subject --gray_subject --save_mask \
--black-threshold 30\
--gray-saturation-threshold 30 --gray-value-threshold 30 \
--edge-grow 2 \
--feather --feather-radius 4 \
--aot-pretrain experiments/G0000000.pt \
--aot-max-size 1000
pip install -r requirements.txt
```
查看已下载模型：
```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
```
## AOT-GAN 修补后端
`--remove-subject` 默认使用 AOT-GAN 修补。
AOT-GAN 依赖 PyTorch（官方仓库测试 Python 3.8 / torch 1.8.1）。建议使用独立虚拟环境或确保兼容版本。
```bash
# 安装依赖（示例）
pip install torch torchvision
```
下载预训练权重后，运行示例：
## 快速开始
处理默认 `images/` 目录中的图片，结果输出到 `output/`：
```bash
python remove_background.py "images/IMG_9259 2.JPG" \
--remove-subject --black-subject --gray-subject --save-mask \
--aot-pretrain experiments/places2.pth
python remove_background.py
```
CPU 无 GPU 时的加速建议（只裁剪主体区域并限制最大边）：
处理单张图片：
```bash
python remove_background.py "images/IMG_9259 2.JPG" \
--remove-subject --black-subject --gray-subject --save-mask \
--aot-pretrain experiments/places2.pth \
--aot-crop --aot-crop-pad 24 --aot-max-size 1400
python remove_background.py input.jpg output.png
```
减少“补脸”倾向：启用随机噪声预填充
处理指定目录：
```bash
python remove_background.py "images/IMG_9259 2.JPG" \
--remove-subject --black-subject --gray-subject --save-mask \
--aot-pretrain experiments/places2.pth \
--aot-crop --aot-crop-pad 64 --aot-max-size 900 \
--aot-noise-prefill --aot-noise-strength 1.0
python remove_background.py my_images/ my_output/
```
## 可调整参数说明
指定为印章场景：
### 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（开关）
- **作用**: 是否启用 alpha matting，提升边缘质量
- **默认**: 关闭（不传 `-a/--alpha-matting`）
- **启用方式**: 传入 `-a` 或 `--alpha-matting`
- **效果**: 有利于细节边缘（毛发/细线），但速度稍慢
#### 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 进行额外后处理
- **默认**: 关闭（不传 `-p/--post-process`）
- **启用方式**: 传入 `-p` 或 `--post-process`
- **效果**: 有助于减少毛边，但可能略损失细节
### 4. 去主体补背景 (remove_subject)
用于“去掉主体并补全背景”。当前仅使用 AOT-GAN 修补。
- **remove_subject（开关）**: 启用去主体补背景（默认关闭，传 `--remove-subject` 开启）
- **aot_root**: AOT-GAN 目录（默认: `AOT-GAN-for-Inpainting`）
- **aot_pretrain**: AOT-GAN 权重文件路径（必填）
- **aot_device**: AOT-GAN 设备（默认: `cpu`）
- **aot_block_num**: AOTBlock 数量（默认: 8）
- **aot_rates**: AOTBlock 膨胀率（默认: `1+2+4+8`）
- **aot_crop（开关）**: 仅对 mask 覆盖区域裁剪修补（默认关闭，传 `--aot-crop` 开启）
- **aot_crop_pad (像素)**: 裁剪边缘留白像素（默认: 0）
- **aot_max_size (像素)**: AOT 输入最大边限制（默认: 0 表示不限制）
- **aot_noise_prefill（开关）**: AOT使用随机噪声预填充（默认关闭）
- **aot_noise_strength (系数)**: 噪声强度（默认: 1.0）
- **mask_dilate (像素)**: mask 膨胀大小（默认: 3）。越大去除范围越大，风险更高
- **mask_blur (像素)**: mask 模糊大小（默认: 3）。越大边缘越柔和但易过度
- **mask_threshold (0-255)**: alpha 阈值（默认: 10）。越大保留越多主体
- **edge_grow (像素)**: 主体边缘额外扩张（默认: 0）。用于清理残留边缘
- **save_mask（开关）**: 保存 mask 方便检查（默认关闭，传 `--save-mask` 开启）
- **black_subject（开关）**: 将黑色内容也视为主体（默认关闭，传 `--black-subject` 开启）
- **black_threshold (0-255)**: 黑色阈值（默认: 50）。越大越容易把浅灰当黑
- **gray_subject（开关）**: 将灰阶内容也视为主体（默认关闭，传 `--gray-subject` 开启）
- **gray_saturation_threshold (0-255)**: 灰阶饱和度阈值（默认: 30）。越大越容易把彩色当灰
- **gray_value_threshold (0-255)**: 灰阶亮度阈值（默认: 200）。越大越容易把浅灰当灰
- **feather（开关）**: 启用边缘过渡（默认关闭，传 `--feather` 开启）
- **feather_radius (像素)**: 过渡半径（默认: 5）。越大过渡越柔和但可能变糊
- **说明**: 过渡仅在 mask 外侧进行，避免把主体边缘带回
### 5. 参数调优建议（针对书画/字迹）
- 先开启 `--remove-subject`，仅看主体遮罩是否覆盖到字迹
- 文字残留：提高 `--black-threshold` 或 `--gray-*` 阈值
- 过度修补：降低 `--black-threshold`、`--gray-value-threshold`，并减小 `--mask-dilate/--mask-blur`
- 边缘不自然：尝试开启 `--feather` 并使用较小的 `--feather-radius`
## 常见问题解决
### 问题1: 前景被过度去除
**解决方案**:
```python
alpha_matting = True
alpha_matting_foreground_threshold = 270 # 增加此值
alpha_matting_background_threshold = 10 # 保持较小
```bash
python remove_background.py input.jpg output.png --artwork-type seal
```
### 问题2: 背景残留太多
**解决方案**:
```python
alpha_matting = True
alpha_matting_foreground_threshold = 240 # 保持默认或减小
alpha_matting_background_threshold = 20 # 增加此值
post_process_mask = True # 启用后处理
强制使用 `rembg`：
```bash
python remove_background.py input.jpg output.png --foreground-mode rembg -m isnet-general-use
```
### 问题3: 边缘不自然、有锯齿
**解决方案**:
```python
alpha_matting = True
alpha_matting_erode_size = 15 # 增加平滑程度
查看完整参数：
```bash
python remove_background.py -h
```
### 问题4: 毛发、头发细节丢失
**解决方案**:
```python
model_name = "birefnet-portrait" # 使用人像专用模型
alpha_matting = True
alpha_matting_foreground_threshold = 270 # 增加以保留细节
alpha_matting_erode_size = 5 # 减小以保留细节
## 输出规则
- 单张图片默认输出为 `*_nobg.png`
- 处理目录时，输出文件写入你指定的输出目录
- 脚本内置的目录批处理只扫描输入目录的第一层文件，不递归子目录
- 如果启用 `--remove-subject`，输出文件名改为 `*_bgfill.<原扩展>` 或 `*_bgfill.jpg`
## 前景提取模式
### `artwork`
默认模式。优先适用于书法、国画、篆刻等纸本图像。
### `auto`
先尝试书画专用掩码；如果结果明显不可信，再回退到 `rembg`。
### `rembg`
强制使用通用抠图模型。适合非书画类图片，或书画专用规则不适合的特殊样本。
## 书画类型
### `auto`
自动兼容书法与印章。
### `calligraphy`
更偏重墨色笔画、灰黑色文字。
### `seal`
更偏重红章、篆刻印记。
## 常用参数
### 书画专用参数
- `--foreground-mode`
- `--artwork-type`
- `--artwork-max-size`
建议：
- 大图先尝试 `--artwork-max-size 1600`
- 超大图如果速度较慢，可降低到 `1200` 或 `1000`
- 红章较多的图片优先试 `--artwork-type seal`
- 纯墨迹优先试 `--artwork-type calligraphy`
### rembg 相关参数
这些参数只在 `--foreground-mode rembg` 或 `auto` 回退到 `rembg` 时生效：
- `-m, --model`
- `-a, --alpha-matting`
- `-ft, --foreground-threshold`
- `-bt, --background-threshold`
- `-es, --erode-size`
- `-p, --post-process`
当前支持的 `rembg` 模型包括：
- `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`
### 去主体补背景参数
启用：
```bash
python remove_background.py input.jpg output.jpg --remove-subject --aot-pretrain experiments/your_model.pt
```
## 推荐配置
常用参数：
### 配置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
- `--aot-root`
- `--aot-pretrain`
- `--aot-device`
- `--aot-block-num`
- `--aot-rates`
- `--aot-crop`
- `--aot-crop-pad`
- `--aot-max-size`
- `--aot-noise-prefill`
- `--aot-noise-strength`
- `--mask-dilate`
- `--mask-blur`
- `--mask-threshold`
- `--edge-grow`
- `--save-mask`
- `--black-subject`
- `--black-threshold`
- `--gray-subject`
- `--gray-saturation-threshold`
- `--gray-value-threshold`
- `--feather`
- `--feather-radius`
一个偏保守的示例：
```bash
python remove_background.py "images/inpaint/IMG_9259 2.JPG" output.jpg \
--remove-subject \
--foreground-mode artwork \
--artwork-type auto \
--aot-pretrain experiments/G0000000.pt \
--aot-crop \
--aot-crop-pad 64 \
--aot-max-size 900 \
--aot-noise-prefill \
--aot-noise-strength 1.0
```
### 配置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
## 典型用法
书法图去背景：
```bash
python remove_background.py input.jpg output.png \
--foreground-mode artwork \
--artwork-type calligraphy
```
### 配置3: 快速处理
```python
model_name = "u2netp"
alpha_matting = False
post_process_mask = False
印章图去背景：
```bash
python remove_background.py input.jpg output.png \
--foreground-mode artwork \
--artwork-type seal
```
## 测试不同参数
通用图片走 `rembg`：
建议按以下顺序调整:
```bash
python remove_background.py input.jpg output.png \
--foreground-mode rembg \
--model birefnet-general
```
1. 先尝试不同的模型
2. 启用alpha_matting
3. 调整foreground_threshold和background_threshold
4. 最后调整erode_size
## 调参建议
每次修改后运行脚本，对比结果。
如果背景没有去掉：
- 先尝试 `--artwork-type calligraphy`
- 再尝试 `--foreground-mode auto`
- 非书画图直接改用 `--foreground-mode rembg`
如果字边仍有底色：
- 先确认原图是否有严重纸纹、阴影或压缩噪声
- 降低 `--artwork-max-size` 可能更快，但通常不利于细节
- 对超大图可以保留 `1600`，必要时单独抽样检查结果
如果去主体补背景不自然：
- 开启 `--aot-crop`
- 增加 `--aot-crop-pad`
- 尝试 `--aot-noise-prefill`
- 减小 `--mask-dilate`、`--mask-blur`
## 校验命令
项目当前可用的检查命令：
```bash
~/venv/bin/python -m pytest -q
~/venv/bin/python -m mypy remove_background.py tests/test_remove_background.py
~/venv/bin/python -m ruff check remove_background.py tests/test_remove_background.py
~/venv/bin/python -m ruff format remove_background.py tests/test_remove_background.py
```
说明：
- `mypy` 与 `ruff` 的项目配置见 [pyproject.toml](/Users/gavin/removeback/pyproject.toml)
- `images/` 与 `output/` 已在静态检查配置里排除
## 已知限制
- CLI 自带的目录批处理不递归子目录
- 书画专用规则仍然可能受极端纸色、重阴影、扫描边框影响
- 某些透明 PNG 在图片预览器里会显示白底或棋盘底，这是预览器合成效果，不代表 alpha 一定有问题
## 仓库结构
```txt
remove_background.py 主脚本
tests/test_remove_background.py
requirements.txt
pyproject.toml
images/ 输入样例
output/ 输出目录
experiments/ AOT-GAN 权重示例
```

BIN
__pycache__/remove_background.cpython-312.pyc
View File

Binary file not shown.

BIN
images/.DS_Store vendored
View File

Binary file not shown.

806
remove_background.py
View File

File diff suppressed because it is too large Load Diff

125
tests/test_remove_background.py
View File

@@ -1,7 +1,9 @@
import importlib
import sys
import tempfile
import types
import unittest
from pathlib import Path
def _import_remove_background():
@@ -18,6 +20,15 @@ class RemoveBackgroundTests(unittest.TestCase):
def setUp(self):
self.mod = _import_remove_background()
def _make_paper_image(self, width=160, height=120):
rng = self.mod.np.random.default_rng(0)
base = self.mod.np.full((height, width, 3), 230, dtype=self.mod.np.uint8)
noise = rng.integers(-8, 9, size=(height, width, 3), dtype=self.mod.np.int16)
image = self.mod.np.clip(base.astype(self.mod.np.int16) + noise, 0, 255).astype(
self.mod.np.uint8
)
return image
def test_alpha_to_mask_threshold(self):
alpha = self.mod.np.array([[0, 9, 10, 255]], dtype=self.mod.np.uint8)
mask = self.mod._alpha_to_mask(alpha, threshold=10)
@@ -67,12 +78,126 @@ class RemoveBackgroundTests(unittest.TestCase):
expanded_min = self.mod._expand_bbox_min_size(bbox, 0, 5, 5, 4)
self.assertEqual(expanded_min, (1, 0, 4, 3))
def test_remove_border_frame_components_keeps_inner_content(self):
mask = self.mod.np.zeros((12, 12), dtype=self.mod.np.uint8)
mask[0, :] = 255
mask[-1, :] = 255
mask[:, 0] = 255
mask[:, -1] = 255
mask[4:8, 5:7] = 255
cleaned = self.mod._remove_border_frame_components(
mask,
min_width_ratio=0.7,
min_height_ratio=0.7,
max_fill_ratio=0.35,
)
self.assertEqual(int(cleaned[0, 0]), 0)
self.assertGreater(cleaned[4:8, 5:7].mean(), 200)
def test_ensure_rgba_size_resizes(self):
img = self.mod.Image.new("RGB", (2, 3), color=(0, 0, 0))
out = self.mod._ensure_rgba_size(img, (4, 5))
self.assertEqual(out.mode, "RGBA")
self.assertEqual(out.size, (4, 5))
def test_extract_artwork_mask_detects_calligraphy_strokes(self):
image = self._make_paper_image()
self.mod.cv2.line(image, (20, 20), (130, 90), (20, 20, 20), thickness=6)
self.mod.cv2.line(image, (45, 15), (45, 105), (30, 30, 30), thickness=5)
mask = self.mod._extract_artwork_mask(
self.mod.Image.fromarray(image),
artwork_type="calligraphy",
max_size=120,
)
self.assertGreater(mask[25:95, 42:49].mean(), 180)
self.assertLess(mask[:15, :15].mean(), 20)
def test_extract_artwork_mask_detects_red_seal(self):
image = self._make_paper_image()
self.mod.cv2.rectangle(image, (40, 30), (120, 95), (165, 30, 40), thickness=5)
self.mod.cv2.line(image, (55, 40), (105, 85), (170, 25, 35), thickness=6)
mask = self.mod._extract_artwork_mask(
self.mod.Image.fromarray(image),
artwork_type="seal",
max_size=120,
)
self.assertGreater(mask[38:92, 38:122].mean(), 40)
self.assertLess(mask[:15, :15].mean(), 20)
def test_extract_artwork_mask_ignores_warm_paper_cast(self):
image = self.mod.np.full(
(120, 160, 3), (222, 188, 150), dtype=self.mod.np.uint8
)
self.mod.cv2.line(image, (30, 20), (30, 100), (28, 28, 28), thickness=6)
self.mod.cv2.line(image, (55, 20), (125, 95), (32, 32, 32), thickness=5)
mask = self.mod._extract_artwork_mask(
self.mod.Image.fromarray(image),
artwork_type="calligraphy",
max_size=120,
)
self.assertGreater(mask[25:100, 26:34].mean(), 180)
self.assertLess(mask[:12, :12].mean(), 20)
def test_mask_to_transparent_image_refines_alpha_inside_coarse_mask(self):
image = self.mod.np.full(
(120, 160, 3), (224, 190, 156), dtype=self.mod.np.uint8
)
self.mod.cv2.line(image, (40, 18), (40, 102), (24, 24, 24), thickness=6)
self.mod.cv2.rectangle(image, (102, 24), (132, 54), (170, 30, 40), thickness=-1)
coarse_mask = self.mod.np.zeros((120, 160), dtype=self.mod.np.uint8)
coarse_mask[10:110, 20:140] = 255
out = self.mod._mask_to_transparent_image(
self.mod.Image.fromarray(image),
coarse_mask,
)
alpha = self.mod.np.array(out.getchannel("A"))
self.assertGreater(alpha[25:95, 36:45].mean(), 110)
self.assertGreater(alpha[30:50, 108:128].mean(), 90)
self.assertLess(alpha[70:90, 70:90].mean(), 25)
def test_remove_background_artwork_mode_skips_rembg(self):
image = self._make_paper_image(width=100, height=80)
self.mod.cv2.line(image, (15, 15), (85, 60), (20, 20, 20), thickness=5)
original_remove = self.mod.remove
def _unexpected_remove(*args, **kwargs):
raise AssertionError("artwork 模式不应调用 rembg")
self.mod.remove = _unexpected_remove
try:
with tempfile.TemporaryDirectory() as tmpdir:
input_path = Path(tmpdir) / "input.png"
output_path = Path(tmpdir) / "output.png"
self.mod.Image.fromarray(image).save(input_path)
self.mod.remove_background(
str(input_path),
str(output_path),
foreground_mode="artwork",
artwork_type="calligraphy",
artwork_max_size=120,
)
out = self.mod.Image.open(output_path)
alpha = self.mod.np.array(out.getchannel("A"))
self.assertEqual(out.mode, "RGBA")
self.assertGreater(alpha[30:55, 35:70].mean(), 70)
self.assertLess(alpha[:10, :10].mean(), 10)
finally:
self.mod.remove = original_remove
if __name__ == "__main__":
unittest.main()