【精华】YOLOv6训练自己的数据集

Github

[美团技术团队]YOLOv6:快速准确的目标检测框架开源

美团AI团队博客

YOLO系列梳理(9)新鲜出炉的第一次品尝YOLOv6

SIoU Loss: More Powerful Learning for Bounding Box Regression

YOLOX: Exceeding YOLO Series in 2021

RepVGG: Making VGG-style ConvNets Great Again

本文主要记录使用情况YOLOv训练自己的数据集， 数据集以Objects365数据集为例.

链接：https://pan.baidu.com/s/1QiWm8hCJus3LstZkz6Mzdw 提取码：wmrx

Objects365数据集为COCO格式数据, 数据集文件格式如下：

Objects365 --Images   --train     --obj365_train_**.jpg   --val     --obj365_val_**.jpg --Annotations   --train     --train.json   --val     --val.json 

YOLOv6默认使用YOLO格式数据集， 中心点坐标采用位置坐标格式， 数据集文件格式如下：

Objects365_yolov6 --images   --train2017     --obj365_train_**.jpg   --val2017     --obj365_val_**.jpg --labels   --train2017     --train2017.txt     --classes.txt     --obj365_train_**.txt     --obj365_train_**.txt   --val2017     --val2017.txt     --classes.txt     --obj365_val_**.txt --obj365_val_**.txt 

通过以下脚本实现COCO格式的数据集转化为YOLO格式的数据集:

#COCO 格式的数据集转化为 YOLO 格式的数据集
# YOLO格式数据集文件结构
'''
--images
  --train2017
    --1_train.jpg
  --val2017
    --2_val.jpg
--labels
  --train2017
    --train2017.txt
    --classes.txt
    --1_train.txt
  --val2017
    --val2017.txt
    --classes.txt
    --2_val.txt
'''

import os
import json
from tqdm import tqdm

def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = box[0] + box[2] / 2.0
    y = box[1] + box[3] / 2.0
    w = box[2]
    h = box[3]
#round函数确定(xmin, ymin, xmax, ymax)的小数位数
    x = round(x * dw, 6)
    w = round(w * dw, 6)
    y = round(y * dh, 6)
    h = round(h * dh, 6)
    return (x, y, w, h)

if __name__ == '__main__':
    # --------------------------------------------------------------------------------------------------------- # 
    json_file = "your/to/path/Objects365/Annotations/train/train.json"    # Objects365 json_path
    ana_txt_save_path = "your/to/path/Objects365_yolov6/labels/train2017"  # anno_txt_save_path
    list_file = open(os.path.join(ana_txt_save_path, 'train2017.txt'), 'w')   
    txt_images_path = 'your/to/path/Objects365_yolov6/images/train2017'
    # --------------------------------------------------------------------------------------------------------- # 

    data = json.load(open(json_file, 'r'))
    if not os.path.exists(ana_txt_save_path):
        os.makedirs(ana_txt_save_path)

    id_map = { 
        } # 数据集的id不连续！重新映射一下再输出！
    with open(os.path.join(ana_txt_save_path, 'classes.txt'), 'w') as f:
        # 写入classes.txt
        for i, category in enumerate(data['categories']):
            f.write(f"{category['name']}\n")
            id_map[category['id']] = i
    # print(id_map)
    #这里需要根据自己的需要，更改写入图像相对路径的文件位置。

    for img in tqdm(data['images']):
        filename = img["file_name"]
        img_width = img["width"]
        img_height = img["height"]
        img_id = img["id"]
        head, tail = os.path.splitext(filename)
        ana_txt_name = head + ".txt"  # 对应的txt名字，与jpg一致
        f_txt = open(os.path.join(ana_txt_save_path, ana_txt_name), 'w')
        for ann in data['annotations']:
            if ann['image_id'] == img_id:
                box = convert((img_width, img_height), ann["bbox"]) f_txt.write("%s %s %s %s %s\n" % (id_map[ann["category_id"]], box[0], box[1], box[2], box[3]))
        f_txt.close()
        #将图片的相对路径写入train2017或val2017的路径
        list_file.write(txt_images_path + '/%s.jpg\n' %(head))
    list_file.close()

模型文件路径: config/yolov6n_objects365.py (新建) 以YOLOv6n为例:

# YOLOv6n model
model = dict(
    type='YOLOv6n',
    pretrained='./weights/yolov6n.pt',
    depth_multiple=0.33,
    width_multiple=0.25,
    backbone=dict(
        type='EfficientRep',
        num_repeats=[1, 6, 12, 18, 6],
        out_channels=[64, 128, 256, 512, 1024],
        ),
    neck=dict(
        type='RepPAN',
        num_repeats=[12, 12, 12, 12],
        out_channels=[256, 128, 128, 256, 256, 512],
        ),
    head=dict(
        type='EffiDeHead',
        in_channels=[128, 256, 512],
        num_layers=3,
        begin_indices=24,
        anchors=1,
        out_indices=[17, 20, 23],
        strides=[8, 16, 32],
        iou_type='ciou'
    )
)

solver = dict(
    optim='SGD',
    lr_scheduler='Cosine',
    lr0=0.00258,
    lrf=0.17,
    momentum=0.779,
    weight_decay=0.00058,
    warmup_epochs=1.33,
    warmup_momentum=0.86,
    warmup_bias_lr=0.0711
)

data_aug = dict(
    hsv_h=0.0188,
    hsv_s=0.704,
    hsv_v=0.36,
    degrees=0.373,
    translate=0.0902,
    scale=0.491,
    shear=0.602,
    flipud=0.00856,
    fliplr=0.5,
    mosaic=1.0,
    mixup=0.243
)

数据集配置文件路径: data/objects365.yaml

针对Objects365此类类别个数超过超过2位数的数据集需修改以下文件:

文件路径: yolov6/data/datasets.py

    @staticmethod
    def check_label_files(args):
        img_path, lb_path = args
        nm, nf, ne, nc, msg = 0, 0, 0, 0, ""  # number (missing, found, empty, message
        try:
            if osp.exists(lb_path):
                nf = 1  # label found
                with open(lb_path, "r") as f:
                    labels = [
                        x.split() for x in f.read().strip().splitlines() if len(x)
                    ]
                    labels = np.array(labels, dtype=np.float32)
                if len(labels):
                    assert all(
                        len(l) == 5 for l in labels
                    ), f"{lb_path}: wrong label format."
                    assert (
                        labels >= 0
                    ).all(), f"{lb_path}: Label values error: all values in label file must > 0"
  # --------------------------------------注释掉本断言-------------------------------------------------------------------- # 
                    # assert (
                    # labels[:, 1:] <= 1
                    # ).all(), f"{lb_path}: Label values error: all coordinates must be normalized"
  # -------------------------------------------------------------------------------------------------------------------- # 

                    _, indices = np.unique(labels, axis=0, return_index=True)
                    if len(indices) < len(labels):  # duplicate row check
                        labels = labels[indices]  # remove duplicates
                        msg += f"WARNING: {lb_path}: {len(labels) - len(indices)} duplicate labels removed"
                    labels = labels.tolist()
                else:
                    ne = 1  # label empty
                    labels = []
            else:
                nm = 1  # label missing
                labels = []

            return img_path, labels, nc, nm, nf, ne, msg
        except Exception as e:
            nc = 1
            msg = f"WARNING: {lb_path}: ignoring invalid labels: {e}"
            return None, None, nc, nm, nf, ne, msg

使用conda或者Docker虚拟环境皆可,此处请自行搭建

• 训练参数配置

  # 文件位置: tools/train.py
  def get_args_parser(add_help=True):
      parser = argparse.ArgumentParser(description='YOLOv6 PyTorch Training', add_help=add_help)
      parser.add_argument('--data-path', default='./data/coco.yaml', type=str, help='path of dataset')
      parser.add_argument('--conf-file', default='./configs/yolov6s.py', type=str, help='experiments description file')
      parser.add_argument('--img-size', type=int, default=640, help='train, val image size (pixels)')
      parser.add_argument('--batch-size', default=32, type=int, help='total batch size for all GPUs')
      parser.add_argument('--epochs', default=400, type=int, help='number of total epochs to run')
      parser.add_argument('--workers', default=4, type=int, help='number of data loading workers (default: 8)')
      parser.add_argument('--device', default='0', type=str, help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
      parser.add_argument('--eval-interval', type=int, default=20, help='evaluate at every interval epochs')
      parser.add_argument('--eval-final-only', action='store_true', help='only evaluate at the final epoch')
      parser.add_argument('--heavy-eval-range', default=50,
                          help='evaluating every epoch for last such epochs (can be jointly used with --eval-interval)')
      parser.add_argument('--check-images', action='store_true', help='check images when initializing datasets')
      parser.add_argument('--check-labels', action='store_true', help='check label files when initializing datasets')
      parser.add_argument('--output-dir', default='./runs/train', type=str, help='path to save outputs')
      parser.add_argument('--name', default='exp', type=str, help='experiment name, saved to output_dir/name')
      parser.add_argument('--dist_url', type=str, default="default url: tcp://127.0.0.1:8888")
      parser.add_argument('--gpu_count', type=int, default=0)
      parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter')
      parser.add_argument('--resume', type=str, default=None, help='resume the corresponding ckpt')
  

• 模型训练

  python tools/train.py --batch-size 32 --conf-file configs/yolov6n_objects365.py --data-path data/objects365.yaml --device 0
  

• 验证参数配置

  # 文件位置: tools/eval.py
  def get_args_parser(add_help=True):
      parser = argparse.ArgumentParser(description='YOLOv6 PyTorch Evalating', add_help=add_help)
      parser.add_argument('--data', type=str, default='./data/coco.yaml', help='dataset.yaml path')
      parser.add_argument('--weights', type=str, default='./weights/yolov6s.pt', help='model.pt path(s)')
      parser.add_argument('--batch-size', type=int, default=32, help='batch size')
      parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
      parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold')
      parser.add_argument('--iou-thres', type=float, default=0.65, help='NMS IoU threshold')
      parser.add_argument('--task', default='val', help='val, or speed')
      parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
      parser.add_argument('--half', default=False, action='store_true', help='whether to use fp16 infer')
      parser.add_argument('--save_dir', type=str, default='runs/val/', help='evaluation save dir')
      parser.add_argument('--name', type=str, default='exp', help='save evaluation results to save_dir/name')
      args = parser.parse_args()
      LOGGER.inf