YOLO - 模型训练入门教程2（样本标注、数据集准备）

作者：hangge | 2024-04-24 08:40

我在前文演示了如何在 Windows 环境下搭建 YOLO 模型的训练环境（点击查看），接下来我将演示如何标注样本并准备用于训练的数据集。

二、样本标注、数据集准备

1，准备样本图片

（1）为方便后续操作与管理，首先我们可以在项目中创建如下层级的空目录：

（2）然后将样本图片放到 own_datas\images\train 文件夹下作为训练集：

提示：这里使用的猫狗图片集可以从百度网盘下载（点击访问，提取码：485q），我从中挑选了部分进行使用（70 张猫、30 张狗）。

2，安装标注工具

（1）在在 Pycharm 终端中执行如下命令安装 labelme 标注工具：

pip install pyqt5 labelme

（2）安装完毕后执行如下命令便会弹出一个窗口，就在这个窗口里进行训练集的标注：

labelme

3，标注图片

（1）开始标注前，我们先点击菜单->文件：

去掉“Save With Image Data”勾选。
点击“更改输出路径”按钮选择标注数据保存到 own_datas\labels\json 目录下。

（2）点击界面上的的 Open Dir（打开目录）按钮，选择 own_datas\images\train 文件夹，就会出现训练集里的图片。右键选择 Create Rectangle（创建矩形），框出图片里的动物：

（3）框选结束后，输入标签名 cat 或者 dog，点击 ok，这个标签就保存下来了。如果有多只动物，就继续框选。

（4）整张图片框选完毕后，点击左侧的保存按钮，根据提示把标注文件保存到路径 own_datas\labels\json 中，文件的格式是 .json。然后点下一幅继续标注，直至所有图片都标注完毕后，该文件夹存放的便是与训练集图片数量相同且对应的 .json 标签文件。

任意打开一个 json 文件，里面内容如下：

{
  "version": "5.4.1",
  "flags": {},
  "shapes": [
    {
      "label": "cat",
      "points": [
        [
          68.6131498470948,
          1.1804281345565792
        ],
        [
          458.8272171253822,
          318.0
        ]
      ],
      "group_id": null,
      "description": "",
      "shape_type": "rectangle",
      "flags": {},
      "mask": null
    }
  ],
  "imagePath": "..\\..\\images\\train\\cat.132.jpg",
  "imageData": null,
  "imageHeight": 336,
  "imageWidth": 499
}

4，转换标签格式

（1）因为 yolov5 只能识别 .txt 格式的标签，还需要把 .json 文件转换成 .txt 文件。我们在项目中新建一个 json2txt.py 文件：

（2）文件内容如下（其中标签名称、输入、输出路径部分根据实际情况进行调整）：

import json
import os

name2id = {'cat': 0, 'dog': 1}  # 标签名称

def convert(img_size, box):
    dw = 1. / (img_size[0])
    dh = 1. / (img_size[1])
    x = (box[0] + box[2]) / 2.0 - 1
    y = (box[1] + box[3]) / 2.0 - 1
    w = box[2] - box[0]
    h = box[3] - box[1]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)

def decode_json(json_floder_path, json_name):
    txt_name = 'D:/PythonProject/yolov5-master/own_datas/labels/txt/' + json_name[0:-5] + '.txt'
    # txt文件夹的绝对路径
    txt_file = open(txt_name, 'w')

    json_path = os.path.join(json_floder_path, json_name)
    data = json.load(open(json_path, 'r', encoding='gb2312', errors='ignore'))

    img_w = data['imageWidth']
    img_h = data['imageHeight']

    for i in data['shapes']:

        label_name = i['label']
        if (i['shape_type'] == 'rectangle'):
            x1 = int(i['points'][0][0])
            y1 = int(i['points'][0][1])
            x2 = int(i['points'][1][0])
            y2 = int(i['points'][1][1])

            bb = (x1, y1, x2, y2)
            bbox = convert((img_w, img_h), bb)
            txt_file.write(str(name2id[label_name]) + " " + " ".join([str(a) for a in bbox]) + '\n')

if __name__ == "__main__":

    json_floder_path = 'D:/PythonProject/yolov5-master/own_datas/labels/json/'
    # json文件夹的绝对路径
    json_names = os.listdir(json_floder_path)
    for json_name in json_names:
        decode_json(json_floder_path, json_name)

（3）运行 json2txt.py，结束后可以在 own_datas\labels\txt 文件夹中看到对应的 .txt 文件：

任意打开一个txt文件，里面内容如下：

0 0.5250501002004008 0.47172619047619047 0.7815631262525049 0.9434523809523809

（4）最后，将 txt 文件夹中的文件全部复制到 own_datas\labels\train 文件夹中。否则，在训练的时候会报错找不到标签。

Python / AI