Samge0 / yolo8-plus-iopaint

本demo使用ultralytics-YOLO8对水印位置进行模型训练&检测，然后使用IOPaint移除检测到的水印。

移除品牌logo水印的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服务：

python iopaint_server.py

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

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

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版

    pip install torch torchvision

  • 【二选一】安装torch-cuda版

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

  • 【必要】安装依赖

    pip install -r requirements.txt

运行方式-demo1：

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

python remove_watermark.py

运行方式-demo2：

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

python remove_watermark_with_onnx.py

错误处理

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

• 点击下载：https://pan.baidu.com/s/1SrxZFkxwpwydn1fuFaWtgw?pwd=6cgb 提取码: 6cgb
• lib文件放到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.x\lib （这里以windows系统为例，其中v11.x是实际安装的cuda版本路径）
• dll文件放到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.x\bin （这里以windows系统为例，其中v11.x是实际安装的cuda版本路径）

相关截图

before	after