【深度学习】使用yolov5对数据进行预标注

介绍

以条形码为例，下载数据集muenster_barcodedb，想使用yolov5对条形码进行检测。muenster_barcodedb数据集有几千张图，没有目标检测标注。

手动标注少数图

随意选95张图出来，手动LabelImg标注，VOC标注格式。

转化VOC到YOLO

填写xmi路径和yolo txt保存路径即可：

# coding:utf-8
import os
import os.path
import xml.etree.ElementTree as ET


def convert_annotation(xmldir: str, txtdir: str, image_id: str, classes: dict):
    in_file = open(os.path.join(xmldir, '%s.xml' % (image_id)), 'r', encoding='UTF-8')
    out_file = open(os.path.join(txtdir, '%s.txt' % (image_id)), 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    size_width = int(size.find('width').text)
    size_height = int(size.find('height').text)
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes[cls]
        xmlbox = obj.find('bndbox')
        b = [float(xmlbox.find('xmin').text),
             float(xmlbox.find('xmax').text),
             float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text)]

        if size_width == 0 or size_height == 0 or b[0] == b[1] or b[2] == b[3]:
            print("不合理的图不再给labels ", image_id)
            # if os.path.exists(xmldir + '%s.xml' % (image_id)):
            # os.remove(xmldir + '%s.xml' % (image_id))
            out_file.close()
            os.remove(os.path.join(txtdir, '%s.txt' % (image_id)))
            return 1

        # 标注越界修正
        if b[0] < 0:
            b[0] = 0
        if b[1] > size_width:
            b[1] = size_width
        if b[2] < 0:
            b[2] = 0
        if b[3] > size_height:
            b[3] = size_height
        txt_data = [round(((b[0] + b[1]) / 2.0 - 1) / size_width, 6),
                    round(((b[2] + b[3]) / 2.0 - 1) / size_height, 6),
                    round((b[1] - b[0]) / size_width, 6),
                    round((b[3] - b[2]) / size_height, 6)]
        if txt_data[0] < 0 or txt_data[1] < 0 or txt_data[2] < 0 or txt_data[3] < 0:
            print("不合理的图不再给labels ", image_id)
            out_file.close()
            os.remove(os.path.join(txtdir, '%s.txt' % (image_id)))
            return 1
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in txt_data]) + '\n')

    in_file.close()
    out_file.close()
    return 0


def listPathAllfiles(dirname):
    result = []
    for maindir, subdir, file_name_list in os.walk(dirname):
        for filename in file_name_list:
            apath = os.path.join(maindir, filename)
            result.append(apath)
    return result


if __name__ == '__main__':
    classes = {
    'barcode': 0}  # 标签名:标签id

    # 给出VOC数据存储地址
    VOCrootpath = r"E:\detection\02barcode\hand"
    xmlpath = os.path.join(VOCrootpath, "xml")
    txtpath = os.path.join(VOCrootpath, "labels")

    if not os.path.exists(xmlpath):
        os.makedirs(xmlpath)
    if not os.path.exists(txtpath):
        os.makedirs(txtpath)

    allfiles = listPathAllfiles(xmlpath)
    print("一共有文件个数:", len(allfiles))
    failNum = 0
    for xmlName in allfiles:
        # xml存储路径，yololabels存储路径，xml文件名称不带.xml后缀，需要的类及其类id的字典
        if convert_annotation(xmlpath, txtpath, os.path.basename(xmlName)[:-4], classes) == 1:
            failNum += 1
    print("失败了多少个文件的labels:", failNum)

转移数据到Linux服务器

训练YOLOV5

添加数据yaml文件

path: /data/dong_xie/datasets/hand  # dataset root dir
train: images
val: images

nc: 1  # number of classes
names: ['barcode']  # class names

选一个合适的权重

训练

batch和epoch少一点，有点担心过拟合：

python train.py --batch-size 16 --data barcode.yaml --img 640 --epochs 50 --weight weights/yolov5m.pt 

Results saved to runs/train/exp23

开服务

yolov5_config.py：

Started server process [30724]
Waiting for application startup.
Application startup complete.
Uvicorn running on http://0.0.0.0:7001 (Press CTRL+C to quit)
127.0.0.1:38726 - "POST /init_process HTTP/1.1" 200
[2022/07/19 22:27:17] layer INFO: algorithm already start 
[2022/07/19 22:27:17] layer INFO: algorithm already start 
algorithm already start 
Fusing layers... 
Model summary: 290 layers, 20852934 parameters, 0 gradients, 47.9 GFLOPs

写labelstxt文件

# -*- coding: utf-8 -*-

import os

import cv2
import numpy as np
import requests


def listPathAllfiles(dirname):
    result = []
    for maindir, subdir, file_name_list in os.walk(dirname):
        for filename in file_name_list:
            apath = os.path.join(maindir, filename)
            result.append(apath)
    return result


if __name__ == '__main__':
    aimpath = r"/data/dong_xie/datasets/Muenster BarcodeDB"  # 填写需要detect的数据在哪
    saveLabelPath = "/data/dong_xie/datasets/bardata"  # 填写存哪里
    if not os.path.exists(os.path.join(saveLabelPath, "images")):
        os.makedirs(os.path.join(saveLabelPath, "images"))
    if not os.path.exists(os.path.join(saveLabelPath, "labels")):
        os.makedirs(os.path.join(saveLabelPath, "labels"))
    cnt = 1
    files = listPathAllfiles(aimpath)
    for fname in files:
        if fname.endswith(".png") or fname.endswith(".jpg"):
            with open(fname, "rb") as fb:
                res = requests.post(url="http://0.0.0.0:7001/analyseImageBytes",
                                    files={
    "file": fb},
                                    data={
    "retain_classes": "0"})

                cnt += 1
                resl = res.json()['data']['detector']
                res = []
                if len(resl):
                    img = cv2.imdecode(np.fromfile(fname, dtype=np.uint8), 1)  # img是矩阵
                    for bbdict in resl:
                        x0, y0, x1, y1 = bbdict['bbox']
                        x = round((x0 + x1) / 2 / img.shape[1], 6)
                        y = round((y0 + y1) / 2 / img.shape[0], 6)
                        w = round((x1 - x0) / img.shape[1], 6)
                        h = round((y1 - y0) / img.shape[0], 6)
                        res.append(" ".join(["2", str(x), str(y), str(w), str(h)]))  # 修改类别id
                    open(os.path.join(saveLabelPath, "labels", "barcode_" + str(cnt).zfill(6) + ".txt"), "w").write(
                        "\n".join(res))  # 修改类别前缀
                    cv2.imencode('.jpg', img)[1].tofile(
                        os.path.join(saveLabelPath, "images", "barcode_" + str(cnt).zfill(6) + ".jpg"))  # 修改类别前缀