【pytorch】yolov5 训练自己的数据集

1.labelme生成标签

2.labelme标签转成yolo标签

# -*- coding: utf-8 -*-  
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import shutil

classes = ["cellphone", "person"]

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

def convert_annotation(path, save, image_id):
    in_file = open("{}/{}.xml".format(path, image_id), "r", encoding="utf-8")
    out_file = open("{}/{}.txt".format(save, image_id), "w", encoding="utf-8")
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find("size")
    w = int(size.find("width").text)
    h = 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.index(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))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

    img_path = os.path.join(path, "{}.jpg".format(image_id))
    save_path = os.path.join(save, "{}.jpg".format(image_id))

    shutil.copy(img_path, save_path)

wd = getcwd()

if __name__ == '__main__':
    # get image_ids
    path = "val_voc"
    save = "val"
    txt_yolo = "val.txt"
    image_ids = []
    for file in listdir(path):
        if file.endswith(".xml"):
            image_id = file.replace(".xml", "")
            image_ids.append(image_id)

    # generate jpg and txt labels
    for image_id in image_ids:
        convert_annotation(path, save, image_id)

    # write yolo
    list_file = open(txt_yolo, "w", encoding="utf-8")
    for image_id in image_ids:
        txt_path = os.path.join(save, "{}.txt".format(image_id))
        img_path = os.path.join(wd, save, "{}.jpg".format(image_id)) # 绝对路径
        lines = open(txt_path, "r", encoding="utf-8")
        write_line = "{} {}\n".format(img_path, " ".join([line.strip() for line in lines]))
        list_file.write(write_line)
    list_file.close()

3.yolov5数据文件格式

── coco128
   ├── images
   │   ├── train
   │   └── val
   └── labels
       ├── train
       └── val

$ cat 00009.txt 
45 0.479492 0.688771 0.955609 0.5955
50 0.637063 0.732938 0.494125 0.510583