基于 YOLO8 + IOPaint 的智能水印移除方案

移除小红书logo水印的 智能演示项目

使用 ultralytics YOLO8 精准检测水印位置,结合 IOPaint 强大的图像修复能力, 实现自动化、高效率的水印移除。支持图片和视频处理,提供多种运行方式。

YOLO8 检测
IOPaint 修复
ONNX 优化
Python 3.8+

核心特性

强大而灵活的水印移除解决方案

智能检测

使用 YOLO8 深度学习模型精准定位小红书水印位置,准确率高,适应性强

高质量修复

基于 IOPaint 的先进图像修复算法,移除水印后保持图像自然清晰

ONNX 加速

支持 ONNX 模型推理,大幅提升检测速度,适合批量处理任务

视频处理

支持视频水印移除,自动处理每一帧并重新合成,保持音视频同步

API 模式

支持 IOPaint API 调用模式,减少批量处理时的初始化开销

灵活部署

支持 Dev Container 快速启动,也支持自定义虚拟环境和远程服务部署

工作流程

简单高效的处理流程

1

加载模型

加载 YOLO8 训练好的水印检测模型 best.pt

2

检测位置

YOLO8 模型识别并定位图片中的水印位置

3

生成遮罩

根据检测结果创建精确的遮罩区域

4

修复输出

IOPaint 基于遮罩进行智能修复,输出干净图像

输入图像

YOLO8 检测

创建遮罩

IOPaint 修复

输出结果

运行方式

三种运行方式,满足不同需求

Demo 1: ultralytics + IOPaint (命令行方式)

使用 PyTorch 模型,适合快速测试

推荐新手
python remove_watermark.py

使用 ultralytics YOLO8 进行水印检测,配合 IOPaint 命令行工具进行修复。脚本:remove_watermark.py

PyTorch 命令行 简单

Demo 2: onnxruntime + IOPaint (命令行方式)

使用 ONNX 模型,推理速度更快

推荐
python remove_watermark_with_onnx.py

使用 ONNX Runtime 进行模型推理,速度更快,内存占用更少。支持 PT 模型自动转换为 ONNX 格式。脚本:remove_watermark_with_onnx.py

ONNX 高性能 批量

Demo 3: 视频水印移除

处理视频文件,逐帧移除水印

实验性
python remove_watermark_video.py

支持视频水印移除,但处理时间较长(4秒测试视频约需291秒)。适合水印位置固定的场景,可复用遮罩优化。脚本:remove_watermark_video.py

视频 逐帧 耗时较长

关键依赖版本

当前开发环境使用的主要依赖

python 3.8.18
torch 2.3.0+cu118
torchvision 0.18.0+cu118
ultralytics 8.2.26
IOPaint 1.3.3
onnxruntime_gpu 1.18.0
onnx 1.16.1
onnx-simplifier 0.4.36
onnxsim 0.4.36
onnxslim 0.1.28

环境配置

快速启动开发环境

Dev Container (推荐)

使用 VS Code 的 Dev Containers 模式,一键启动完整开发环境。

查看配置说明

手动配置 (可选)

安装 PyTorch (二选一)

pip install torch torchvision
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

安装依赖

pip install -r requirements.txt

常见问题

解答使用过程中的疑问

如何启用 IOPaint API 模式?

remove_watermark_with_onnx.pyremove_watermark.py 中设置 USE_IOPAINT_API=True,然后启动 IOPaint 服务:python iopaint_server.py

如何配置自定义的 IOPaint 服务地址?

可以通过环境变量 IOPAINT_SERVER_HOST 配置,或在 iopaint_api_utils.py 中直接修改 IOPAINT_SERVER_HOST 常量。

遇到 zlibwapi.dll 错误怎么办?

需要下载相关的 DLL 文件并放置到 CUDA 目录:

  • lib 文件放到 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.x\lib
  • dll 文件放到 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.x\bin

百度网盘下载:提取码 6cgb

视频处理为什么这么慢?

4秒测试视频(143帧)使用 IOPaint 批量移除水印耗时约280秒。如果水印位置固定,可以优化为使用单一遮罩而非为每帧创建遮罩,大幅减少处理时间。

处理效果对比

水印移除前后对比

处理前

原始图片(带水印)

处理后

干净图片(水印已移除)

完整对比视频和截图请查看项目 GitHub 页面