本文主要记录使用YOLOv6训练自己数据集的过程, 数据集以Objects365数据集为例.
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格式的数据集:
'''
--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]
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"
ana_txt_save_path = "your/to/path/Objects365_yolov6/labels/train2017"
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 = {}
with open(os.path.join(ana_txt_save_path, 'classes.txt'), 'w') as f:
for i, category in enumerate(data['categories']):
f.write(f"{category['name']}\n")
id_map[category['id']] = i
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"
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()
list_file.write(txt_images_path + '/%s.jpg\n' %(head))
list_file.close()
模型文件路径: config/yolov6n_objects365.py (新建) 以YOLOv6n为例:
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
train: your/to/path/Objects365_yolov6/images/train2017
val: your/to/path/Objects365_yolov6/images/val2017
test: your/to/path/Objects365_yolov6/images/val2017
nc: 365
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"
]
针对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, ""
try:
if osp.exists(lb_path):
nf = 1
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"
_, indices = np.unique(labels, axis=0, return_index=True)
if len(indices) < len(labels):
labels = labels[indices]
msg += f"WARNING: {lb_path}: {len(labels) - len(indices)} duplicate labels removed"
labels = labels.tolist()
else:
ne = 1
labels = []
else:
nm = 1
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虚拟环境皆可,此处请自行搭建
- 训练参数配置
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
- 验证参数配置
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
- 推理参数配置
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
Original: https://blog.csdn.net/qq_36722887/article/details/125749759
Author: LeeZhao@
Title: 【精华】YOLOv6训练自己的数据集
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/629594/
转载文章受原作者版权保护。转载请注明原作者出处！