YOLOV8训练自己的数据集图文实战,包含voc数据集处理代码

yolov8官方链接: link
本文章是以labelimg标注好的voc数据集为基础,通过转换格式训练模型,

一,安装

pip install ultralytics

pip install git+https://github.com/ultralytics/ultralytics.git@main

link

二,数据集准备

数据集格式如下

├── ultralytics
└── datasets
└── VOCData
└── images
└── img1.jpg…
└── Annotations
└── img1.xml…

在这里插入图片描述

2.1,同文件目录下运行以下代码,生成label文件夹,并生成训练集和测试集


# coding:utf-8import os
import random
import argparseparser = argparse.ArgumentParser()
#xml文件的地址，根据自己的数据进行修改 xml一般存放在Annotations下
parser.add_argument('--xml_path', default='Annotations', type=str, help='input xml label path')
#数据集的划分，地址选择自己数据下的ImageSets/Main
parser.add_argument('--txt_path', default='labels', type=str, help='output txt label path')
opt = parser.parse_args()trainval_percent = 1.0
train_percent = 0.9
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)
if not os.path.exists(txtsavepath):os.makedirs(txtsavepath)num = len(total_xml)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')for i in list_index:name = total_xml[i][:-4] + '\n'if i in trainval:file_trainval.write(name)if i in train:file_train.write(name)else:file_val.write(name)else:file_test.write(name)file_trainval.close()
file_train.close()
file_val.close()
file_test.close()

2.2接下来运行以下代码,修改标签名称

# -*- coding: utf-8 -*-import xml.etree.ElementTree as ET
from tqdm import tqdm
import os
from os import getcwdsets = ['train', 'val', 'test']classes = ['f','fd','fe']   # 改成自己的类别def convert(size, box):dw = 1. / (size[0])dh = 1. / (size[1])x = (box[0] + box[1]) / 2.0 - 1y = (box[2] + box[3]) / 2.0 - 1w = box[1] - box[0]h = box[3] - box[2]x = x * dww = w * dwy = y * dhh = h * dhreturn x, y, w, hdef convert_annotation(image_id):# try:in_file = open('VOCData/Annotations/%s.xml' % (image_id), encoding='utf-8')out_file = open('VOCData/labels/%s.txt' % (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').textcls = obj.find('name').textif cls not in classes or int(difficult) == 1:continuecls_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))b1, b2, b3, b4 = b# 标注越界修正if b2 > w:b2 = wif b4 > h:b4 = hb = (b1, b2, b3, b4)bb = convert((w, h), b)out_file.write(str(cls_id) + " " +" ".join([str(a) for a in bb]) + '\n')# except Exception as e:#     print(e, image_id)wd = getcwd()
for image_set in sets:if not os.path.exists('VOCData/labels/'):os.makedirs('VOCData/labels/')image_ids = open('VOCData/labels/%s.txt' %(image_set)).read().strip().split()list_file = open('VOCData/%s.txt' % (image_set), 'w')for image_id in tqdm(image_ids):list_file.write('VOCData/JPEGImages/%s.jpg\n' % (image_id))convert_annotation(image_id)list_file.close()

在这里插入图片描述

2.3数据集转移到制定目录中

import os
import shutiltrain_file = 'VOCData/images/train/'if not os.path.exists(train_file):os.makedirs(train_file)
with open('VOCData/labels/train.txt', 'r') as f:counts = f.readlines()for i in counts:imgname=i.split('\n')[0]shutil.move('VOCData/images/'+imgname+'.jpg', train_file)val_file = 'VOCData/images/val/'if not os.path.exists(val_file):os.makedirs(val_file)
with open('VOCData/labels/val.txt', 'r') as f:counts = f.readlines()for i in counts:imgname=i.split('\n')[0]shutil.move('VOCData/images/'+imgname+'.jpg', val_file)xml_train_file = 'VOCData/labels/train/'if not os.path.exists(xml_train_file):os.makedirs(xml_train_file)
with open('VOCData/labels/train.txt', 'r') as f:counts = f.readlines()for i in counts:imgname=i.split('\n')[0]shutil.move('VOCData/labels/'+imgname+'.txt', xml_train_file)xml_val_file = 'VOCData/labels/val/'if not os.path.exists(xml_val_file):os.makedirs(xml_val_file)
with open('VOCData/labels/val.txt', 'r') as f:counts = f.readlines()for i in counts:imgname=i.split('\n')[0]shutil.move('VOCData/labels/'+imgname+'.txt', xml_val_file)

在这里插入图片描述

三,修改配置文件,开始训练模型

link

vi  ultralytics/ultralytics/cfg/datasets/coco8.yaml

3.1修改路径

在这里插入图片描述

3.2训练模型


from ultralytics import YOLO# Load a model
model = YOLO("yolov8l.pt")  # load a pretrained model (recommended for training)# Train the model with 2 GPUs
results = model.train(data="/home/ai-developer/ultralytics/my_coco8.yaml", epochs=50, imgsz=640, batch=32,device=[0, 1],lr0=0.0001)

在这里插入图片描述

3.推理


from ultralytics import YOLO
import os
# 加载模型
model = YOLO("yolov8m.yaml")  # 从头开始构建新模型
model = YOLO("/ultralytics/runs/detect/train18/weights/best.pt")  # 加载预训练模型（建议用于训练）imgfile=./test_images'
save_path='./results'# 使用模型
for air,file,filelist in os.walk(imgfile):for filename in filelist:imgname=os.path.join(air,filename)results = model(imgname)  # 对图像进行预测for result in results:boxes = result.boxes  # Boxes object for bounding box outputsresult.show()  # display to screensavename=os.path.join(save_path,filename)result.save(savename)  # save to disk