记录下labelImg标注数据到YOLOv8训练的过程,其中容易遇到labelImg的坑
数据集处理
首先在mydata下创建4个文件夹
images文件夹下存放着所有的图片,包括训练集和测试集等。后续会根据代码进行划分。
json文件夹里存放的是labelImg标注的所有数据。需要注意的是,json文件的命名应与images文件夹中的图片一一对应。labels文件夹是空的,后续会根据代码将json转化为YOLOv8支持的训练数据集。接下来需要创建一个 split_train_val.py
文件,放在mydata目录下,用于将images文件夹中的图片划分为训练集和测试集。代码如下:
import os
import json
import random
import argparseclass DatasetSplitter:def __init__(self, json_path, txt_path):self.json_path = json_pathself.txt_path = txt_pathself.trainval_percent = 1.0self.train_percent = 0.9self.total_json = os.listdir(json_path)self.num = len(self.total_json)self.list_index = list(range(self.num))if not os.path.exists(txt_path):os.makedirs(txt_path)def split_dataset(self):tv = int(self.num * self.trainval_percent)tr = int(tv * self.train_percent)trainval = random.sample(self.list_index, tv)train = random.sample(trainval, tr)file_trainval = open(os.path.join(self.txt_path, 'trainval.txt'), 'w')file_test = open(os.path.join(self.txt_path, 'test.txt'), 'w')file_train = open(os.path.join(self.txt_path, 'train.txt'), 'w')file_val = open(os.path.join(self.txt_path, 'val.txt'), 'w')for i in self.list_index:name = self.total_json[i][:-5] + '\n' # Assuming filenames end with '.json'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()if __name__ == "__main__":parser = argparse.ArgumentParser()parser.add_argument('--json_path', default='json', type=str, help='input json label path')parser.add_argument('--txt_path', default='dataSet', type=str, help='output txt label path')opt = parser.parse_args()dataset_splitter = DatasetSplitter(opt.json_path, opt.txt_path)dataset_splitter.split_dataset()
运行后会在dataSet生成四个文件.
然后前往mydata目录的上一级目录,例如我这里是test
目录下,创建voc_label.py
文件.用于将labelImg标注的json数据转化为txt文本数据.
请注意,YULO的标注框中的x和y表示矩形框的左上角,而labelImg中的(x,y,w,h)可能表示矩形框的中心点坐标。在使用时需要确认,以免训练出来的YOLO存在偏差。
可以通过以下代码draw_picture_by_json.py
进行绘图测试,以下我是默认labelImg的格式是(x_center,y_center,w,h)的格式进行绘制.
from typing import Dict
from PIL import Image, ImageDraw
import jsondef draw_pic_by_raw_data(image_path:str, output_path:str, json_data: Dict):data = json_dataimg = Image.open(image_path)draw = ImageDraw.Draw(img)# Load JSON file with bounding box information# Iterate through bounding boxes and draw them on the imageannotations = data[0]["annotations"]for bbox in annotations:label = bbox["label"]x_center, y_center, width, height = (bbox["coordinates"]["x"],bbox["coordinates"]["y"],bbox["coordinates"]["width"],bbox["coordinates"]["height"])#FIXME TODO 从中心坐标x,y LabelImg 计算左上角坐标x, y = x_center - width / 2, y_center - height / 2# Draw bounding boxdraw.rectangle([x, y, x+width, y+height], outline="red", width=2)# Draw labeldraw.text((x, y-15), label, fill="red")# Save the resultimg.save(output_path)img.show()if __name__ == "__main__":test_img_path = "3157.jpg"test_json_path = "3157.json"output_path = "out.jpg"with open(test_json_path, 'r') as f:data = json.load(f)draw_pic_by_raw_data(test_img_path, output_path, data)
通过以上代码,如果矩形框正确,则验证了labelImg的x、y、w和h坐标应该是x_center、y_center、w和h的情况。需要在将json转换为txt文本时进行处理。这里给出voc_label.py
的代码如下:请注意修改classes中的类别,修改为和mydata.yaml的顺序一致,否则会导致标签错误。
代码运行前请在
mydata
目录的上一级,也就是我这里的test
目录下创建一个paper_data
文件夹,用于等会在mydata.yaml中指定路径.
import os
import json
from os import getcwdsets = ['train', 'val', 'test']
classes = ["A", "B", "C", "D", "E"] # 请根据您的实际类别进行修改 对齐yaml文件中的names
abs_path = os.getcwd()
print(abs_path)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):in_file = open('mydata/json/%s.json' % (image_id), 'r', encoding='UTF-8')out_file = open('mydata/labels/%s.txt' % (image_id), 'w')data = json.load(in_file)for obj in data[0]["annotations"]:w = obj["coordinates"]["width"]h = obj["coordinates"]["height"]difficult = 0 # JSON数据中没有difficult字段,设为0cls = obj["label"]if cls not in classes or difficult == 1:continuecls_id = classes.index(cls)##注意:这里的box是左上角和右下角的坐标,而不是中心点坐标x_center = obj["coordinates"]["x"]y_center = obj["coordinates"]["y"]width = obj["coordinates"]["width"]height = obj["coordinates"]["height"]x, y = x_center - width / 2, y_center - height / 2b = x, x + width, y, y + height# b = obj["coordinates"]["x"], obj["coordinates"]["x"] + obj["coordinates"]["width"], \# obj["coordinates"]["y"], obj["coordinates"]["y"] + obj["coordinates"]["height"]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')wd = getcwd()
for image_set in sets:if not os.path.exists('mydata/labels/'):os.makedirs('mydata/labels/')image_ids = open('mydata/dataSet/%s.txt' % (image_set)).read().strip().split()list_file = open('paper_data/%s.txt' % (image_set), 'w')for image_id in image_ids:list_file.write(abs_path + '/mydata/images/%s.jpg\n' % (image_id))convert_annotation(image_id)list_file.close()
yaml文件配置
截止到这里数据preprocess处理完毕.接下来进行训练前的yaml文件配置.
修改为刚才创建的paper_data的绝对路径
并进行names的配置,注意顺序和数量需要与
voc_label.py
里面的列表顺序一致.
最后修改yolov8.yaml文件配置
yolov8.yaml的路径一般在ultralytics/ultralytics/cfg/models/v8
目录下
修改nc为names中的类别数量
训练模型
yolo task=detect mode=train model=yolov8s.yaml data=mydata.yaml epochs=25 batch=16
这些模型可以支持多种尺寸,包括小(n)、中(s)、中大(m)、大(l)和超大(x)。模型的尺寸参数越大,所需的显存和训练时间也越多。要切换模型尺寸,只需在"model=yolov8"后面加上相应的尺寸参数(n、s、m、l、x),然后再加上.yaml即可。另外,epochs设置为25,batch大小为16。epochs表示训练的轮数, batch为每个批次的样本数量为16.
推理
训练结束后可以通过以下代码进行测试.
yolo predict model=xxxx/weights/best.pt source=xxxx/mydata/testvedio.mp4 imgsz=1280