小红书 Logo 水印检测 Demo

核心特性

基于 YOLOv8 的智能水印检测方案，支持自定义数据集训练

精确检测

基于 YOLOv8 深度学习模型，对小红书 Logo 水印进行高精度定位识别

简单标注

使用 LabelImg 进行可视化标注，支持 YOLO 格式导出，快速构建训练数据集

自定义训练

支持自定义数据集进行训练，灵活配置模型参数，满足不同场景需求

工作流程

从数据标注到模型部署的完整流程

1

数据准备

收集样本图片，放置到 datasets/data/images 目录

2

数据标注

使用 LabelImg 进行水印位置标注，导出 YOLO 格式

3

数据处理

运行 Process.py 将数据集拆分为训练、测试、验证集

4

模型训练

执行 train.py 训练模型，test.py 进行推理验证

环境配置

快速搭建开发环境，支持多种部署方式

Dev Containers 【推荐】

使用 VSCode 的 Dev Containers 模式，参考 .devcontainer/README.md 进行配置

cp .devcontainer/docker-compose-demo.yml .devcontainer/docker-compose.yml

主要关注 /root/.cache 路径映射、代理 PROXY 配置

虚拟环境【可选】

【二选一】安装 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

关键依赖版本

python==3.8.18
torch==2.3.0+cu118
torchvision==0.18.0+cu118
ultralytics==8.2.28
labelImg==1.8.6

训练与推理

简单的命令即可完成模型训练和推理验证

模型训练

python train.py

执行训练脚本，开始对标注数据进行模型训练

模型推理

python test.py

对测试图片进行水印检测，输出预测结果

自定义数据集训练

使用自己的数据进行模型训练，满足特定需求

1

安装 labelImg

pip install labelImg

2

启动 labelImg

labelImg

3

准备数据

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

4

数据标注

• 在 labelImg 中打开 datasets/data/images 图集
• 保存格式选择 YOLO
• 建议点击 File -> YOLO 保存全局默认格式

5

数据集拆分

cd datasets && python Process.py

将 images 目录拆分为 train、test、val 目录

6

执行训练

• 执行 python train.py 进行训练
• 执行 python test.py 进行推理

常见问题

关于项目使用和配置的常见疑问

如何提高检测准确率？

增加训练样本数量、提高标注质量、调整模型参数（如 imgsz、epochs等）可以有效提高检测准确率。建议至少使用 100 张以上标注图片进行训练。

支持 GPU 加速吗？

支持。安装 CUDA 版本的 torch 可以启用 GPU 加速，大幅提升训练和推理速度。参考环境配置中的 torch-cuda 版安装步骤。

如何自动移除检测到的水印？

可以结合 IOPaint 工具实现自动水印移除。具体使用方法请参考 yolo8-plus-iopaint 仓库。

训练时出现内存不足错误怎么办？

可以减小 batch size、降低图片分辨率（imgsz）、减少数据集规模。如果使用 CPU 训练，建议使用较小的模型配置。

可以使用自己的图片数据吗？

可以。参考"自定义数据集训练"部分，将您的图片放到指定目录，使用 labelImg 进行标注，然后按照标准流程进行训练即可。