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

---

文章目录

---

(1) YOLOv6

​ Github

​ [美团技术团队] YOLOv6：又快又准的目标检测框架开源啦

​ 美团 AI 团队博客

(2) 参考博客

​ YOLO 系列梳理（九）初尝新鲜出炉的 YOLOv6

​ SIoU Loss: More Powerful Learning for Bounding Box Regression

​ YOLOX: Exceeding YOLO Series in 2021

​ RepVGG: Making VGG-style ConvNets Great Again

​ 本文主要记录使用 YOLOv6 训练自己数据集的过程，数据集以 Objects365 数据集为例.

2 数据集获取

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

3 数据集转化

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()

4 工程文件配置

(1) 配置模型文件

模型文件路径: 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
)

(2) 配置数据集文件

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

train: your/to/path/Objects365_yolov6/images/train2017
val: your/to/path/Objects365_yolov6/images/val2017
test: your/to/path/Objects365_yolov6/images/val2017
#anno_path: your/to/path/Objects365_yolov6/annotations/instances_val2017.json   # 该标签为程序自动生成，不能指定为原COCO数据集的标签文件，路径只需修改根目录即可
# number of classes
nc: 365

# class names
names: [
    "human","sneakers","chair","hat","lamp","bottle","cabinet/shelf","cup","car","glasses","picture/frame","desk","handbag",
    "street lights","book","plate","helmet","leather shoes","pillow","glove","potted plant","bracelet","flower","monitor",
    "storage box","plants pot/vase","bench","wine glass","boots","dining table","umbrella","boat","flag","speaker",
    "trash bin/can","stool","backpack","sofa","belt","carpet","basket","towel/napkin","slippers","bowl","barrel/bucket",
    "coffee table","suv","toy","tie","bed","traffic light","pen/pencil","microphone","sandals","canned","necklace",
    "mirror","faucet","bicycle","bread","high heels","ring","van","watch","combine with bowl","sink","horse","fish",
    "apple","traffic sign","camera","candle","stuffed animal","cake","motorbike/motorcycle","wild bird","laptop",
    "knife","cellphone","paddle","truck","cow","power outlet","clock","drum","fork","bus","hanger","nightstand",
    "pot/pan","sheep","guitar","traffic cone","tea pot","keyboard","tripod","hockey stick","fan","dog","spoon",
    "blackboard/whiteboard","balloon","air conditioner","cymbal","mouse","telephone","pickup truck","orange","banana",
    "airplane","luggage","skis","soccer","trolley","oven","remote","combine with glove","paper towel","refrigerator",
    "train","tomato","machinery vehicle","tent","shampoo/shower gel","head phone","lantern","donut","cleaning products",
    "sailboat","tangerine","pizza","kite","computer box","elephant","toiletries","gas stove","broccoli","toilet","stroller",
    "shovel","baseball bat","microwave","skateboard","surfboard","surveillance camera","gun","Life saver","cat","lemon",
    "liquid soap","zebra","duck","sports car","giraffe","pumpkin","Accordion/keyboard/piano","radiator","converter",
    "tissue","carrot","washing machine","vent","cookies","cutting/chopping board","tennis racket","candy",
    "skating and skiing shoes","scissors","folder","baseball","strawberry","bow tie","pigeon","pepper","coffee machine",
    "bathtub","snowboard","suitcase","grapes","ladder","pear","american football","basketball","potato","paint brush",
    "printer","billiards","fire hydrant","goose","projector","sausage","fire extinguisher","extension cord","facial mask",
    "tennis ball","chopsticks","Electronic stove and gas stove","pie","frisbee","kettle","hamburger","golf club","cucumber",
    "clutch","blender","tong","slide","hot dog","toothbrush","facial cleanser","mango","deer","egg","violin","marker",
    "ship","chicken","onion","ice cream","tape","wheelchair","plum","bar soap","scale","watermelon","cabbage","router/modem",
    "golf ball","pine apple","crane","fire truck","peach","cello","notepaper","tricycle","toaster","helicopter","green beans",
    "brush","carriage","cigar","earphone","penguin","hurdle","swing","radio","CD","parking meter","swan","garlic","french fries",
    "horn","avocado","saxophone","trumpet","sandwich","cue","kiwi fruit","bear","fishing rod","cherry","tablet","green vegetables",
    "nuts","corn","key","screwdriver","globe","broom","pliers","hammer","volleyball","eggplant","trophy","board eraser","dates",
    "rice","tape measure/ruler","dumbbell","hamimelon","stapler","camel","lettuce","goldfish","meat balls","medal","toothpaste",
    "antelope","shrimp","rickshaw","trombone","pomegranate","coconut","jellyfish","mushroom","calculator","treadmill","butterfly",
    "egg tart","cheese","pomelo","pig","race car","rice cooker","tuba","crosswalk sign","papaya","hair dryer","green onion","chips",
    "dolphin","sushi","urinal","donkey","electric drill","spring rolls","tortoise/turtle","parrot","flute","measuring cup","shark",
    "steak","poker card","binoculars","llama","radish","noodles","mop","yak","crab","microscope","barbell","Bread/bun","baozi",
    "lion","red cabbage","polar bear","lighter","mangosteen","seal","comb","eraser","pitaya","scallop","pencil case","saw",
    "table tennis  paddle","okra","starfish","monkey","eagle","durian","rabbit","game board","french horn","ambulance","asparagus",
    "hoverboard","pasta","target","hotair balloon","chainsaw","lobster","iron","flashlight"
]

(3) 其他文件修改

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

文件路径: YOLOv6/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

5 模型训练及推理

(1) 训练环境搭建

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

(2) 模型训练

• 训练参数配置

  # 文件位置: 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
  

(3) 模型验证

• 验证参数配置

  # 文件位置: 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.info(args)
      return args
  

• python tools/eval.py --data data/objects365.yaml  --batch-size 32 --weights weights/yolov6n.pt --task val
  

(4) 模型推理

• 推理参数配置

  # 文件位置: tools/infer.py
  def get_args_parser(add_help=True):
      parser = argparse.ArgumentParser(description='YOLOv6 PyTorch Inference.', add_help=add_help)
      parser.add_argument('--weights', type=str, default='weights/yolov6s.pt', help='model path(s) for inference.')
      parser.add_argument('--source', type=str, default='data/images', help='the source path, e.g. image-file/dir.')
      parser.add_argument('--yaml', type=str, default='data/coco.yaml', help='data yaml file.')
      parser.add_argument('--img-size', type=int, default=640, help='the image-size(h,w) in inference size.')
      parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold for inference.')
      parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold for inference.')
      parser.add_argument('--max-det', type=int, default=1000, help='maximal inferences per image.')
      parser.add_argument('--device', default='0', help='device to run our model i.e. 0 or 0,1,2,3 or cpu.')
      parser.add_argument('--save-txt', action='store_true', help='save results to *.txt.')
      parser.add_argument('--save-img', action='store_false', help='save visuallized inference results.')
      parser.add_argument('--classes', nargs='+', type=int, help='filter by classes, e.g. --classes 0, or --classes 0 2 3.')
      parser.adkd_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS.')
      parser.add_argument('--project', default='runs/inference', help='save inference results to project/name.')
      parser.add_argument('--name', default='exp', help='save inference results to project/name.')
      parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels.')
      parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences.')
      parser.add_argument('--half', action='store_true', help='whether to use FP16 half-precision inference.')
  
  args = parser.parse_args()
      LOGGER.info(args)
      return args
  

• python tools/infer.py --weights weights/yolov6n.pt --yaml data/objects365.yaml --source your/to/images.jpg