跳转至

整体架构

项目结构

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 环境变量加载

GPU 可选依赖([gpu] extra)

用途
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)