整体架构
项目结构
src/remove_ai_watermarks/
├── cli.py # Click CLI 入口
├── gemini_engine.py # 可见水印移除(纯 CPU)
├── invisible_engine.py # 不可见水印移除(需要 GPU)
├── metadata.py # AI 元数据检测/移除
├── humanizer.py # 模拟胶片噪声注入
├── face_protector.py # YOLO 人脸检测与保护
├── assets/ # 内嵌 alpha 背景图
│ ├── gemini_bg_48.png # 48×48 Gemini 水印 alpha map
│ └── gemini_bg_96.png # 96×96 Gemini 水印 alpha map
└── noai/ # 核心水印去除子系统
├── watermark_remover.py # 主协调器
├── img2img_runner.py # img2img 管线执行
├── watermark_profiles.py # 模型配置与强度预设
├── c2pa.py # C2PA/JUMBF 解析
├── cleaner.py # AI 元数据清除
├── extractor.py # 元数据提取
├── isobmff.py # AVIF/HEIF/JXL 容器级 C2PA 移除
├── constants.py # 共享常量
├── utils.py # 格式检测工具
├── progress.py # 终端进度动画
└── ctrlregen/ # CtrlRegen 子管线 (ICLR 2025)
├── engine.py # 端到端编排
├── pipeline.py # 自定义 SD ControlNet Img2Img
├── ip_adapter.py # DINOv2 IP-Adapter
├── color.py # 色彩匹配后处理
└── tiling.py # 大图分块处理
模块依赖图
graph TD
CLI[cli.py] --> GE[gemini_engine.py]
CLI --> IE[invisible_engine.py]
CLI --> META[metadata.py]
GE --> ASSETS[assets/alpha maps]
IE --> WM[noai/watermark_remover.py]
IE --> FP[face_protector.py]
IE --> HUM[humanizer.py]
WM --> I2I[noai/img2img_runner.py]
WM --> PROFILES[noai/watermark_profiles.py]
WM --> CR[noai/ctrlregen/engine.py]
WM --> META2[noai/cleaner.py]
CR --> PIPE[ctrlregen/pipeline.py]
CR --> IPA[ctrlregen/ip_adapter.py]
CR --> CLR[ctrlregen/color.py]
CR --> TIL[ctrlregen/tiling.py]
META2 --> C2PA[noai/c2pa.py]
META2 --> ISOBMFF[noai/isobmff.py]
FP --> YOLO[ultralytics YOLO]
style CLI fill:#4a9eff,color:#fff
style GE fill:#ff6b6b,color:#fff
style IE fill:#51cf66,color:#fff
style META fill:#ffd43b,color:#333
CLI 命令体系
基于 Click 框架,入口点定义在 pyproject.toml:
[project.scripts]
remove-ai-watermarks = "remove_ai_watermarks.cli:main"
| 命令 |
功能 |
需要 GPU |
visible |
可见 Gemini 水印移除(离线) |
否 |
invisible |
不可见水印移除(SynthID 等) |
是 |
metadata |
AI 元数据检测与剥离 |
否 |
all |
全流程:visible → invisible → metadata |
是(invisible 阶段) |
batch |
批量处理目录 |
取决于 mode |
依赖分层
核心依赖(始终安装)
| 包 |
用途 |
pillow |
图像 I/O |
piexif |
EXIF 读写 |
numpy |
数组运算 |
opencv-python-headless |
图像处理、NCC 检测、inpainting |
click |
CLI 框架 |
rich |
终端进度条 |
python-dotenv |
环境变量加载 |
| 包 |
用途 |
torch |
PyTorch 推理引擎 |
diffusers |
HuggingFace 扩散管线 |
transformers |
DINOv2/CLIP 编码器 |
accelerate |
设备自动调度 |
controlnet-aux |
Canny 边缘检测器 |
safetensors |
模型权重加载 |
ultralytics |
YOLOv8 人脸检测 |
color-matcher |
色彩校正 |
运行入口
CLI
# 全流程
remove-ai-watermarks all image.png -o clean.png
# 单独命令
remove-ai-watermarks visible image.png -o clean.png
remove-ai-watermarks invisible image.png -o clean.png --humanize 4.0
remove-ai-watermarks metadata image.png --check
remove-ai-watermarks metadata image.png --remove
Python API
from remove_ai_watermarks.gemini_engine import GeminiEngine
from remove_ai_watermarks.invisible_engine import InvisibleEngine
from remove_ai_watermarks.metadata import has_ai_metadata, remove_ai_metadata
# 可见水印
engine = GeminiEngine()
result = engine.detect_watermark(image)
clean = engine.remove_watermark(image)
# 不可见水印
inv = InvisibleEngine()
inv.remove_watermark("input.png", "output.png")
# 元数据
if has_ai_metadata(path):
remove_ai_metadata(path, output_path)