LOGO识别及数据集训练

背景介绍

近年来，随着数字媒体的迅猛发展，涌现出了大量优秀的自媒体创作者。然而博主在创作过程中，为了美化图片的显示质量，恶意删除和篡改图片自带的logo（水印），严重侵害了原创者的权益。因此，如何使用AI算法识别出违禁图片，进而辅助人们对原创者的权益进行保护。调研发现，YOLO (You Only Look Once) 是一个流行的目标检测算法，能够实现图像分类、图像分割、目标跟踪以及姿态估计等。因此，本推文展示了获取训练数据集—数据标注—模型训练—结果后处理的全过程，具体内容如下所示：

图1 yolo v8算法的应用场景.

项目框架

水印检测能够帮助企业追踪图片的来源，辅助工作人员对图像进行分类。近期，跟企业合作开展了logo水印识别，后续采用lama-cleaner对图片进行修复，还原了图片原始信息。整个项目的框架为：1.AI识别出目标的种类，输出目标在图片中的位置；2.制作mask蒙版；3.图片修复，具体的执行框架如下所示：

环境搭建

调研发现，图像识别算法有非常多种，主要包含：CNN（需要大量标注）、R-CNN（精度好，但是速度相对慢）、GAN（图像生成，训练复杂）、YOLO（速度快，对小物体，重叠物体识别效果差）、HOG（行人检测，对复杂背景适应性差）以及U-Net（主要用于医学图像分割）等。项目从图像识别的准确率，训练时间成本，历史经验积累以及难度等方面综合考虑，最终决定采用了yolo算法进行目标识别。其中，环境的搭建过程如下所示：

export http_proxy=http://192.168.0.177:7890
export https_proxy=http://192.168.0.177:7890

git clone https://github.com/ultralytics/ultralytics
cd ultralytics
pip install -e .
 
sudo yolo predict model=yolov8m.pt source='https://ultralytics.com/images/bus.jpg' save=True

数据标注

执行目标检测前，我们需要给模型喂一定的数据，得到训练好的权重文件。其中，YOLO算法需要的数据格式为：2 0.793750 0.878125 0.334375 0.168750；其中第一列为类别，第二第三列为矩形框中心点的相对位置，第三第四列为矩形框的长宽。附1：数据集制作过程中，要尽可能地包含所有数据的特征，保证结果的准确可靠性；附2：数据标注过程中要尽可能的准确，避免漏标，误标以及多标等情况；项目中采用了labeling软件进行数据标注，具体过程如下：

pip install labelimg
输入“labelimg”，打开图像标注工具主界面（找到安装位置或者添加环境变量）C:\Users\znmay\AppData\Roaming\Python\Python39\Scripts

图2 数据标注过程

模型训练

本例程中不再赘述模型训练过程，感兴趣的读者可以参考网上的教程，工作过程中可以使用下述命令行进行模型训练和图片识别，具体为：

# 模型训练
sudo yolo detect train data=datasets/mytrain.yaml model=yolov8m.yaml epochs=300 batch=4
# 图片识别 
sudo yolo detect val data=datasets/zmlogo/mytrain.yaml model=runs/detect/train7/weights/best.pt batch=4
# 断点训练
sudo yolo detect train data=datasets/mytrain.yaml model=weights/last.pt epochs=300 batch=4 save=True resume=True

图3 训练过程：采用3090显卡进行训练.

附1：模型训练的控制参数非常多，具体有：task=detect, mode=train, model=yolov8m.yaml, data=datasets/mytrain.yaml, epochs=300, patience=50, batch=16, imgsz=640, save=True, save_period=-1, cache=False, device=None, workers=8, project=None, name=None, exist_ok=False, pretrained=True, optimizer=auto, verbose=True, seed=0, deterministic=True, single_cls=False, rect=False, cos_lr=False, close_mosaic=10, resume=False, amp=True, fraction=1.0, profile=False, freeze=None, overlap_mask=True, mask_ratio=4, dropout=0.0, val=True, split=val, save_json=False, save_hybrid=False, conf=None, iou=0.7, max_det=300, half=False, dnn=False, plots=True, source=None, show=False, save_txt=False, save_conf=False, save_crop=False, show_labels=True, show_conf=True, vid_stride=1, stream_buffer=False, line_width=None, visualize=False, augment=False, agnostic_nms=False, classes=None, retina_masks=False, boxes=True, format=torchscript, keras=False, optimize=False, int8=False, dynamic=False, simplify=False, opset=None, workspace=4, nms=False, lr0=0.01, lrf=0.01, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, warmup_bias_lr=0.1, box=7.5, cls=0.5, dfl=1.5, pose=12.0, kobj=1.0, label_smoothing=0.0, nbs=64, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, mosaic=1.0, mixup=0.0, copy_paste=0.0, cfg=None, tracker=botsort.yaml, save_dir=ultralytics/runs/detect/train5

目前觉得比较有用的参数有：save_period：checkpoint的保存频率，默认不保存；imgsz：图像尺寸；optimizer：迭代方法；resume：断电等意外情况时，继续训练；lr0：初始学习速率。

附2：后续有机会的话整一个好的显卡（意向的有3090、4090或者H100二手显卡，A100），对模型训练过程加速非常非常明显；采用显卡仅需要1个小时的计算任务，采用纯CPU可能需要20多小时，等的太急人了。

后处理

后续我们把文件路径，图像尺寸，检测的类别以及矩形框的坐标等信息输出到TXT文件中，具体的过程如下所示：

图4 logo识别结果展示

附1：整个过程中，想要达到非常好的效果大约需要500-1000张数据集，这块还是很磨人的；

附2：不管啥事情，测试Demo一般都能够比较容易的实现，在工程应用中想要提高一点点就需要付出相对多的时间。