使用yolo训练品牌logo的水印检测模型的demo

, , , , , ,

前言

本demo使用labelImg对26张样本图片的水印位置进行标注，ultralytics-YOLO8对水印位置进行模型训练&检测。

如果需要使用 ultralytics-YOLO8 + IOPaint 进行组合，自动移除yolo识别的目标水印，请点击yolo8-plus-iopaint仓库查看。

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

1
2
3
4
5
6
7
python==3.8.18
torch==2.3.0+cu118
torchvision==0.18.0+cu118
ultralytics==8.2.28

# labelImg is used to label the training data
labelImg==1.8.6

环境配置

  • 【推荐】使用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

训练

1
python train.py

推理

1
python test.py

自定义数据集进行训练：

  • 安装labelImg

    1
    pip install labelImg

  • 启动labelImg

    1
    labelImg

  • 清理或备份旧的数据集，将需要训练的新数据图集放到datasets/data/images目录，参与训练的图片宽高最好一致，训练前需要在train.py中配置imgsz图片宽高信息

  • 在labelImg打开datasets/data/images的图集进行标注，保存格式选择YOLO（建议点击File -> YOLO保存全局默认YOLO导出后，重新打开labelImg，可在后续保存标注时避免频繁切换导出格式）

  • 标注完毕后，执行命令将datasets/data/images拆分为datasets/data/train、datasets/data/test、datasets/data/val

    1
    cd datasets && python Process.py

  • 按前面文档所示，执行python train.py进行训练，执行python test.py进行推理

相关截图

  • labelImg标注界面
    labelImg

  • 训练后的模型预测结果
    output-result