Yolov5 训练自己的数据集 - 识别螺丝

一. 准备自己的数据集

准备好自己的所需要的图片，使用labelme对图片进行标注

1. labelme的安装

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple labelme

使用上述命令安装labelme，安装源使用清华源。

2. labelme使用介绍

步骤1：在命令行执行labelme命令打开软件

步骤2：打开图片文件夹

步骤3：点击edit选择图形对目标进行标注，在这里我使用圆进行标注

步骤4：输入目标的标签，可用数字代替

步骤5：保存

步骤6：点击下一张图片

3.labelme生成的json文件介绍

4.处理数据集

按照yolov5的数据集目录结构存放图片和标签数据：

datasets
|--images
|--labels

images：存放图片

labels：存放标注的label数据

yolov5的label数据格式是txt文件，且需要满足一定的格式，如下所示：

0 0.284268465909091 0.4578302556818182 0.2331680030387468 0.2914600037984335
|         |____________________|               |                   |
类别              中心点坐标                   矩形宽               矩形高

基于以上格式，所以需要将label生成的json文件转换成以上格式的txt文件。

另外由于我标注的是圆，但yolov5是按照矩形框进行训练的，所以还需要将圆形转换成矩形，根据圆的半径画一个圆的外切正方形。

转换的代码如下：

import json
import os
import math

file_names = os.listdir('./')
#遍历目录下所有文件
for filename in file_names:
    #处理json文件
    if ".json" in filename:
        #读取文件
        file = open(filename, 'rb')
        data = json.loads(file.read())
        shapes = data["shapes"][0]
        label = shapes["label"]
        point0 = shapes["points"][0]
        point1 = shapes["points"][1]
        file.close()

        #获取两个点的坐标
        x1 = point0[0]
        y1 = point0[1]
        x2 = point1[0]
        y2 = point1[1]
        
        #txt文件名保持和图片名相同
        newfilename_1 = filename[0:-5]
        newfilename_2 = ".txt"
        print(newfilename_1+newfilename_2)
        #保存到../labels/目录
        textfile = open("../labels/"+newfilename_1+newfilename_2, 'w')
        
        #计算半径
        c_r = math.sqrt(pow((x2-x1),2)+pow(y2-y1,2))
        c_r = 2*c_r #直径
        #拼接数据，并做归一化处理，图片大小640x512
        data_str = str(label) + " " + str(x1/640.0) + " " + str(y1/512.0) + " " + str(c_r/640.0) + " " + str(c_r/512.0)
        print(label, x1/640.0, y1/512.0,c_r/640.0,c_r/512.0)
        #保存文件
        textfile.write(data_str)
        textfile.close()
        

以上就得到了一个完整的数据集

二. 修改数据集配置文件

可以直接复制yolov5/data/coco128.yaml文件，并修改内容，如下所示：

# YOLOv5  by Ultralytics, GPL-3.0 license
# COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
# Example usage: python train.py --data mydata.yaml
# parent
# ├── yolov5
# └── datasets
#     └── images
#     |—— labels


# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets  # dataset root dir
train: images  # train images (relative to 'path') 128 images
val: images  # val images (relative to 'path') 128 images
test:  # test images (optional)

# Classes
names:
  0: luosi

因为我只有一个类别，所以names下只有一种。

三. 训练数据集

python3 train.py --img 640 --batch 16 --epochs 5 --data data/mydata.yaml --cfg models/yolov5s.yaml

四. 使用自己的训练的模型进行推理

python3 detect.py --weights runs/train/exp6/weights/best.pt --source ../N_609.jpg --data data/mydata.yaml

训练推理两步具体可以参考Yolov5训练数据集-入门篇-CSDN博客

推理结果如下：

图片资源https://download.csdn.net/download/u010492025/88529074