Watermark segmentation

Source | HN Comments

该项目是 DiffusionDynamics.ai 和 clear.photo 合作开发的开源水印分割技术，基于深度学习和 PyTorch。它使用合成数据增强，精确检测 logo 和文本水印。代码参考了相关研究，提供了一个极简、易于理解的基线，用于生成水印掩码。项目包含训练和推理工具，支持在 Apple M 系列芯片和 NVIDIA GPU 上进行训练。核心是数据生成策略，通过随机化水印实现模型鲁棒性。最终输出水印区域的分割掩码，可用于后续水印移除。

水印分割技术 (Watermark Segmentation)

Diffusion-Dynamics / watermark-segmentation Public

由 DiffusionDynamics.ai 和 clear.photo 共同开发的开源水印分割技术。运用深度学习和基于 PyTorch 的合成数据增强技术，精确检测 logo/文本。该极简代码参考了顶级研究成果，实现了强大且适应性强的水印移除方案。

License

MIT license

1 star 0 forks Branches Tags Activity

Diffusion-Dynamics/watermark-segmentation

main

Branches Tags

文件夹和文件

Name | Name | Last commit message | Last commit date ---|---|---|--- logos | logos .gitignore | .gitignore LICENSE | LICENSE README.md | README.md dataset.py | dataset.py readmeimage-2.png | readmeimage-2.png readmeimage.png | readmeimage.png requirements.txt | requirements.txt watermark-segmentation.ipynb | watermark-segmentation.ipynb

View all files

历史

2 Commits

Watermark Segmentation

该仓库由 Diffusion Dynamics 提供，展示了我们首款产品 clear.photo 背后的核心水印分割技术。这项工作借鉴了图像恢复任务中关于扩散模型的研究见解。

引言

有效的水印移除依赖于准确识别水印在图像中的精确位置和形状。此代码解决了第一个关键步骤：水印分割。

我们提出了一种深度学习方法，该方法经过训练可以生成突出显示水印区域的掩码。本仓库侧重于分割基于 logo 的水印，展示了一种适用于各种水印类型的强大技术。所采用的方法借鉴了图像分割领域的进展。

本仓库旨在整合近期可见水印移除和分割研究中的关键思想，包括以下文献中提出的技术：

arXiv:2108.03581 "Visible Watermark Removal via Self-calibrated Localization and Background Refinement"
arXiv:2012.07616 "WDNet: Watermark-Decomposition Network for Visible Watermark Removal"
arXiv:2312.14383 "Removing Interference and Recovering Content Imaginatively for Visible Watermark Removal"
arXiv:2502.02676 "Blind Visible Watermark Removal with Morphological Dilation"

它将这些概念提炼为一个最小、功能性的代码库，纯粹专注于分割任务。目标是提供一个清晰、易于理解的基线，易于修改和构建，甚至允许在带有 Apple M 系列芯片的笔记本电脑等消费级硬件上进行微调。它可以作为一个基础示例，展示适用于构建更复杂工具（如 clear.photo）的核心技术。

背景：分割在水印移除中的作用

一个典型的先进水印移除流程包括：

分割： 生成一个精确的掩码，将水印像素与背景图像内容隔离。
修复/恢复： 使用掩码来引导算法（通常是生成式的，如扩散模型）智能地“填充”先前被水印占据的区域，无缝地重建底层图像。

可视化效果

代码库概述

本项目提供了训练水印分割模型并将其用于推理的必要工具。关键组件包括：

watermark-segmentation.ipynb： 包含端到端工作流程的 Jupyter notebook：
- 模型定义（利用 segmentation_models.pytorch 库）。
- 使用 dataset.py 设置数据加载器。
- 使用 pytorch-lightning 实现训练循环。
- 应用训练好的模型到新图像的推理逻辑。
dataset.py： 一个 Python 脚本，定义了负责生成训练数据的 Dataset 类。它动态地将 logo 水印（来自 logos/ 目录）应用到具有随机属性（比例、旋转、不透明度、位置、混合模式）的背景图像上，以创建多样化的训练样本及其对应的 ground truth 掩码。这种数据增强策略是模型鲁棒性的关键。
requirements.txt： 列出了所有必要的 Python 依赖项。
*.pth： 来自不同训练 epoch 的模型权重。
logos/： 一个用于示例水印 logo 的目录。使用与您的用例相关的 logo 填充此目录。
lightning_logs/： PyTorch Lightning 保存训练日志和检查点的默认目录。

快速开始

按照以下步骤设置和运行项目：

1. 准备工作：

Python 3.10 或更高版本。
wget 和 unzip（或用于下载和提取数据集的等效工具）。

python --version
# 确保显示 3.10.x 或更高版本

2. 克隆仓库：

git clone https://github.com/Diffusion-Dynamics/watermark-segmentation
cd watermark-segmentation

3. 安装依赖：

建议使用虚拟环境：

python -m venv venv
source venv/bin/activate # On Windows use `virtual\Scripts\activate`
pip install -r requirements.txt

4. 准备数据集：

背景图像： 您需要一个多样化的背景图像数据集。原始设置使用 Flickr8k。下载并将其提取（或您选择的数据集）到一个已知位置。

# 使用 Flickr8k 的示例
wget https://github.com/jbrownlee/Datasets/releases/download/Flickr8k/Flickr8k_Dataset.zip
unzip Flickr8k_Dataset.zip
# 确保 notebook 指向 'Flicker8k_Dataset' 目录

水印 Logo： 将您所需的水印 logo 图像（建议使用具有透明度的 PNG 格式）放入 logos/ 目录中。 dataset.py 脚本将随机使用这些图像。

5. 运行 Notebook：

启动 Jupyter Lab 或 Jupyter Notebook：

jupyter lab

打开 watermark-segmentation.ipynb。
至关重要的是，如果路径与默认路径不同，请更新 notebook 单元格中指向您的背景图像数据集和 logo 目录的路径。
执行单元格以训练模型并执行推理。

模型架构和训练

核心分割模型利用了 segmentation-models-pytorch 库提供的标准分割架构。该库提供了具有各种 backbone 的预实现模型。为了加速训练并减少对大量训练数据的需求，该模型利用了在 ImageNet 数据集上预训练的 backbone。

训练使用 pytorch-lightning 进行管理，从而简化了训练循环，便于多 GPU 训练并集成了日志记录。训练过程与 Apple M 系列芯片（通过 MPS）和 NVIDIA GPU（通过 CUDA）兼容。虽然 NVIDIA GPU 通常提供更快的训练速度，但在 Apple M 系列芯片上进行微调是可行的，并且通常可以在几个小时内完成。

实现良好性能和泛化的关键在于 dataset.py 中的数据生成策略。这涉及合成数据增强：在训练期间将多样化和随机化的水印（改变大小、位置、不透明度、旋转、混合模式）动态地应用于干净的背景图像上。这迫使模型学习在各种条件下识别水印，使其对未见过的水印和背景具有鲁棒性。这种方法受到相关图像恢复研究（[arXiv:2502.02676](https://github.com/Diffusion-Dynamics/https:/arxiv.org/abs/2502.02676））中讨论的技术的启发。

用法 (推理)

该 notebook (watermark-segmentation.ipynb) 演示了如何：

加载包含水印的输入图像。
预处理用于模型的图像。
运行模型以获得分割掩码（通常是概率图）。
后处理掩码（例如，阈值化）以获得二值分割。一个常见的改进步骤，受到“Blind Visible Watermark Removal with Morphological Dilation”等作品中讨论的形态学膨胀等技术的启发（[arXiv:2502.02676](https://github.com/Diffusion-Dynamics/https:/arxiv.org/abs/2502.02676）），可以在这里应用，以潜在地稍微扩展掩码，确保为后续移除步骤提供更好的覆盖，尽管这篇特定的论文侧重于移除过程本身内的改进阶段。

输出掩码精确地识别了水印区域，准备好用于下游移除任务。