yolov5数据集的制作、划分-以识别高空作业及安全带佩戴为例

一、数据

1、数据集

天池算法赛——广东电网智慧现场作业挑战赛 赛道三：识别高空作业及安全带佩戴

地址：天池赛-广东电网智慧现场作业挑战赛-赛道三：识别高空作业及安全带佩戴_数据集-飞桨AI Studio星河社区

这个数据集3_images.tar是图片，an_rname.csv是对应的标签信息，

但yolo的数据集格式是

而labels里面是每张图片对应的标签信息，我们要将csv文件转换成每一张图片名对应的txt文件

2、csv 转 json

参考：天池算法赛——广东电网智慧现场作业挑战赛 赛道三：识别高空作业及安全带佩戴_天池赛道三_小屋*的博客-CSDN博客

'''
官方给出的csv中的
{
 "meta":{},
 "id":"88eb919f-6f12-486d-9223-cd0c4b581dbf",
 "items":
[
     {"meta":{"rectStartPointerXY":[622,2728],"pointRatio":0.5,"geometry":[622,2728,745,3368],"type":"BBOX"},"id":"e520a291-bbf7-4032-92c6-dc84a1fc864e","properties":{"create_time":1620610883573,"accept_meta":{},"mark_by":"LABEL","is_system_map":false},"labels":{"鏍囩":"ground"}}
     {"meta":{"pointRatio":0.5,"geometry":[402.87,621.81,909,1472.01],"type":"BBOX"},"id":"2c097366-fbb3-4f9d-b5bb-286e70970eba","properties":{"create_time":1620610907831,"accept_meta":{},"mark_by":"LABEL","is_system_map":false},"labels":{"鏍囩":"safebelt"}}
     {"meta":{"rectStartPointerXY":[692,1063],"pointRatio":0.5,"geometry":[697.02,1063,1224,1761],"type":"BBOX"},"id":"8981c722-79e8-4ae8-a3a3-ae451300d625","properties":{"create_time":1620610943766,"accept_meta":{},"mark_by":"LABEL","is_system_map":false},"labels":{"鏍囩":"offground"}}
 ],
 "properties":{"seq":"1714"},"labels":{"invalid":"false"},"timestamp":1620644812068
 }
'''

import pandas as pd
import json
import os
from PIL import Image

df = pd.read_csv("an_rname.csv",header=None)
df_img_path = df[4]
df_img_mark = df[5]
# print(df_img_mark)
# 统计一下类别,并且重新生成原数据集标注文件，保存到json文件中
dict_class = {
    "badge": 0,
    "offground": 0,
    "ground": 0,
    "safebelt": 0
}
dict_lable = {
    "badge": 1,
    "offground": 2,
    "ground": 3,
    "safebelt": 4
}
data_dict_json = []
image_width, image_height = 0, 0
ids = 0
false = False  # 将其中false字段转化为布尔值False
true = True  # 将其中true字段转化为布尔值True
for img_id, one_img in enumerate(df_img_mark):
    # print('img_id',img_id)
    one_img = eval(one_img)["items"]
    # print('one_img',one_img)
    one_img_name = df_img_path[img_id]
    img = Image.open(os.path.join("./", one_img_name))
    # print(os.path.join("./", one_img_name))
    ids = ids + 1
    w, h = img.size
    image_width += w
    # print(image_width)
    image_height += h
    # print(one_img_name)
    i=1
    for one_mark in one_img:
        # print('%d      '%i,one_mark)

        one_label = one_mark["labels"]['标签']
        # print('%d      '%i,one_label)
        try:
            dict_class[str(one_label)] += 1
            # category = str(one_label)
            category = dict_lable[str(one_label)]
            bbox = one_mark["meta"]["geometry"]
        except:
            dict_class["badge"] += 1  # 标签为"监护袖章(红only)"表示类别"badge"
            # category = "badge"
            category = 1
            bbox = one_mark["meta"]["geometry"]
        i+=1

        one_dict = {}
        one_dict["name"] = str(one_img_name)
        one_dict["category"] = category
        one_dict["bbox"] = bbox
        data_dict_json.append(one_dict)
print(image_height / ids, image_width / ids)
print(dict_class)
print(len(data_dict_json))
print(data_dict_json[0])
with open("./data.json2", 'w') as fp:
    json.dump(data_dict_json, fp, indent=1, separators=(',', ': '))  # 缩进设置为1，元素之间用逗号隔开 ， key和内容之间 用冒号隔开
    fp.close()

得到json文件，但还是所有图片的json文件

3、json 转 txt

将json文件yolov5可用的标签格式

json_to_txt.py 这里将所有的标签都减了一，可以不改，自己对的上就可以，当前标签：“badge”: 0,“offground”: 1,“ground”: 2,“safebelt”:3 bbox做了归一化（这个分数据集，有的数据集格式不一样，具体情况具体改）

可以读取images文件夹中的多种格式的图片

import json
import os
import cv2

file_name_list = {}

with open("./data.json2", 'r', encoding='utf-8') as fr:
    data_list = json.load(fr)

label = 0
[x1, y1, x2, y2] = [0, 0, 0, 0]

for data_dict in data_list:
    for k, v in data_dict.items():
        if k == "category":
            label = v
        if k == "bbox":
            [x1, y1, x2, y2] = v
        if k == "name":
            file_name, file_extension = os.path.splitext(v)
            file_name = file_name[9:]

    if not os.path.exists('./data1/'):
        os.mkdir('./data1/')

    image_path = os.path.join('./3_images/', file_name + file_extension)
    
    # Check if the file exists
    if not os.path.exists(image_path):
        print(f"Image not found: {image_path}")
        continue

    img = cv2.imread(image_path)

    if img is None:
        print(f"Failed to read image: {image_path}")
        continue
    
    size = img.shape  # (h, w, channel)
    dh = 1.0 / size[0]
    dw = 1.0 / size[1]
    x = (x1 + x2) / 2.0
    y = (y1 + y2) / 2.0
    w = x2 - x1
    h = y2 - y1
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh

    content = f"{label-1} {x} {y} {w} {h}\n"

    with open(os.path.join('./data1/', file_name + '.txt'), 'a+', encoding='utf-8') as fw:
        fw.write(content)

运行代码生成了yolo可用的标签信息

4、数据集划分

参考yolov5训练自己的数据集（一文搞定训练）_yolov5训练数据_小屋*的博客-CSDN博客

但说的有点复杂，不一定要对应的train.txt文件

可以用这种格式

二、训练

1、准备yaml文件

在data文件夹中复制一个coco.yaml文件，重命名为你的数据集名字，修改地址，nc，names

2、修改参数

在train.py中找到参数位置，修改合适的参数

一些参数的注释

def parse_opt(known=False):
    parser = argparse.ArgumentParser()
    #指定权重，''置空，从头开始训练自己的
    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
    #model路径里面5个yaml，depth_multiple:乘64，width_multiple乘3不同
    parser.add_argument('--cfg', type=str, default='models\####.yaml', help='model.yaml path')
    #model路径里面选择一个数据集格式的yaml文件
    parser.add_argument('--data', type=str, default='data\####.yaml', help='dataset.yaml path')
    #data\hyps
    parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml', help='hyperparameters path')
    #迭代次数0-299
    parser.add_argument('--epochs', type=int, default=1)
    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs, -1 for autobatch')
    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
    #矩阵推理。尺寸可以改，节省
    parser.add_argument('--rect', action='store_true', help='rectangular training')
    ##指定意外断了的位置，继续训练
    parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
    parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
    #只在最后一轮测试 
    parser.add_argument('--noval', action='store_true', help='only validate final epoch')
    parser.add_argument('--noautoanchor', action='store_true', help='disable AutoAnchor')
    #不保存绘图
    parser.add_argument('--noplots', action='store_true', help='save no plot files')
    #寻找参数
    parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
    parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
    parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"')
    #样本不平衡，训练不好的图片在下一轮增加一些权重
    parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')

    parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
    #指定单多类别
    parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
    parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer')
    parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
    #多线程
    parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
    parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name')
    
    parser.add_argument('--name', default='exp', help='save to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    #图像大于640开启效果好 
    parser.add_argument('--quad', action='store_true', help='quad dataloader')
    parser.add_argument('--cos-lr', action='store_true', help='cosine LR scheduler')
    #免疫标签打错的情况
    parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
    #
    parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')
    #冻结，迁移学习
    parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone=10, first3=0 1 2')
    parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
    #随机数种子玄学数值，经验
    parser.add_argument('--seed', type=int, default=0, help='Global training seed')
    parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')

    # Weights & Biases arguments
    parser.add_argument('--entity', default=None, help='W&B: Entity')
    parser.add_argument('--upload_dataset', nargs='?', const=True, default=False, help='W&B: Upload data, "val" option')
    parser.add_argument('--bbox_interval', type=int, default=-1, help='W&B: Set bounding-box image logging interval')
    parser.add_argument('--artifact_alias', type=str, default='latest', help='W&B: Version of dataset artifact to use')

    return parser.parse_known_args()[0] if known else parser.parse_args()