强大而灵活的水印移除解决方案
使用 YOLO8 深度学习模型精准定位小红书水印位置,准确率高,适应性强
基于 IOPaint 的先进图像修复算法,移除水印后保持图像自然清晰
支持 ONNX 模型推理,大幅提升检测速度,适合批量处理任务
支持视频水印移除,自动处理每一帧并重新合成,保持音视频同步
支持 IOPaint API 调用模式,减少批量处理时的初始化开销
支持 Dev Container 快速启动,也支持自定义虚拟环境和远程服务部署
简单高效的处理流程
加载 YOLO8 训练好的水印检测模型 best.pt
YOLO8 模型识别并定位图片中的水印位置
根据检测结果创建精确的遮罩区域
IOPaint 基于遮罩进行智能修复,输出干净图像
输入图像
YOLO8 检测
创建遮罩
IOPaint 修复
输出结果
三种运行方式,满足不同需求
使用 PyTorch 模型,适合快速测试
python remove_watermark.py
使用 ultralytics YOLO8 进行水印检测,配合 IOPaint 命令行工具进行修复。脚本:remove_watermark.py
使用 ONNX 模型,推理速度更快
python remove_watermark_with_onnx.py
使用 ONNX Runtime 进行模型推理,速度更快,内存占用更少。支持 PT 模型自动转换为 ONNX 格式。脚本:remove_watermark_with_onnx.py
处理视频文件,逐帧移除水印
python remove_watermark_video.py
支持视频水印移除,但处理时间较长(4秒测试视频约需291秒)。适合水印位置固定的场景,可复用遮罩优化。脚本:remove_watermark_video.py
当前开发环境使用的主要依赖
快速启动开发环境
安装 PyTorch (二选一)
pip install torch torchvision
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
安装依赖
pip install -r requirements.txt
解答使用过程中的疑问
在 remove_watermark_with_onnx.py 或 remove_watermark.py 中设置 USE_IOPAINT_API=True,然后启动 IOPaint 服务:python iopaint_server.py
可以通过环境变量 IOPAINT_SERVER_HOST 配置,或在 iopaint_api_utils.py 中直接修改 IOPAINT_SERVER_HOST 常量。
需要下载相关的 DLL 文件并放置到 CUDA 目录:
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.x\libC:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.x\bin百度网盘下载:提取码 6cgb
4秒测试视频(143帧)使用 IOPaint 批量移除水印耗时约280秒。如果水印位置固定,可以优化为使用单一遮罩而非为每帧创建遮罩,大幅减少处理时间。