YOLO+IOPaint组合，自动识别并移除水印的demo

, , , , , ,

前言

本demo使用ultralytics-YOLO8对水印位置进行检测，然后使用IOPaint移除yolo识别的目标水印。

本demo中使用的last.pt模型来自yolo8-watermark-brand仓库。

本demo支持使用iopaint的api方式去除水印，只需在remove_watermark_with_onnx.py 或 remove_watermark.py中配置USE_IOPAINT_API=True，可减少批量操作时iopaint命令行方式的初始化耗时。

如果配置USE_IOPAINT_API=True，需要先启动iopaint服务：

1
python iopaint_server.py

当然，也可以选择对接单独部署的iopaint服务，只需要在iopaint_api_utils.py中配置自定义的IOPAINT_SERVER_HOST即可。

当前开发环境使用的关键依赖版本

1
2
3
4
5
6
7
8
9
10
11
12
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

# the onnx dependency is to automatically export the onnx model at train time
onnx==1.16.1
onnx-simplifier==0.4.36
onnxsim==0.4.36
onnxslim==0.1.28

环境配置

  • 【推荐】使用vscode的Dev Containers模式，参考.devcontainer/README.md

  • 【可选】其他虚拟环境方式

    • 【二选一】安装torch-cpu版 
      1
      pip install torch torchvision
    • 【二选一】安装torch-cuda版 
      1
      pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
    • 【必要】安装依赖 
      1
      pip install -r requirements.txt

运行方式-demo1：

ultralytics + IOPaint（命令行方式），脚本：remove_watermark.py

1
python remove_watermark.py

运行方式-demo2：

onnxruntime + IOPaint（命令行方式），脚本：remove_watermark_with_onnx.py
pt转onnx模型可参考yolo_utils.py的mian函数

1
python remove_watermark_with_onnx.py

错误处理

1、如果遇到Could not locate zlibwapi.dll. Please make sure it is in your library path错误，需要下载相关dll放置到目标位置：

相关截图

before after
before after