【万字长文】深度学习2 yolov5修改为自己的数据集

数据预处理

使用labelme可以直接导出适用于yolo模型的txt文本数据，也可以直接导出默认的json数据结构，后面我会提供代码进行转换。自行进行标注，图片与标注一一对应，更多要求不赘述。因为我做最简单的检索模型，不做切割，所以用矩形框进行框选就可以。
我的数据集大致如下，老师给的真实数据不能外泄，里面就是一些标注位置及类型：

格式转换：

1、新建目录

标注好的数据放在原始数据集original_datasest里面，注意英文路径下。新建data目录。

2、新建下面四个包存放数据

Annotations：存放转换之后图片的xml文件，images：存放数据集的图片文件，ImageSets ：存放数据集按比例分割之后自动生成的train.txt，val.txt，test.txt和trainval.txt四个文件，labels：存放xml文件转换后的txt标准格式标签。

3、json转xml

import os
import xml.domimport numpy as np
import codecs
import json
import glob
import cv2
import shutil# 1.标签路径
labelme_path = r"D:\something\homework\deeplearning\Mydataset\original_dataset\Thyroid"  # 原始json、bmp标注数据路径,需要更换成自己的数据集名称
saved_path = r"D:\something\homework\deeplearning\Mydataset\data\Annotations"  # 保存路径# 2.获取待处理文件
files = glob.glob("%s/*.json" % (labelme_path))# 3.读取标注信息并写入 xml
for json_filename in files:json_file = json.load(open(json_filename, "r", encoding="utf-8"))i = 0# 图像名字，若图像格式不是jpg，需要修改此处img_name = json_filename.replace(".json", ".jpg")height, width, channels = cv2.imread(img_name).shape# xml名字xmlName = os.path.join(saved_path, json_filename.split("\\")[-1].replace(".json", ".xml"))with codecs.open(xmlName, "w", "utf-8") as xml:print(2)xml.write('<annotation>\n')xml.write('\t<folder>' + 'jpg' + '</folder>\n')xml.write('\t<filename>' + img_name + '</filename>\n')# -------------------------------------------------xml.write('\t<source>\n')xml.write('\t\t<database>hulan</database>\n')# --------------------------------------------------xml.write('\t</source>\n')# -----------------------------------------------------------xml.write('\t<size>\n')xml.write('\t\t<width>' + str(width) + '</width>\n')xml.write('\t\t<height>' + str(height) + '</height>\n')xml.write('\t\t<depth>' + str(channels) + '</depth>\n')# ------------------------------------------------xml.write('\t</size>\n')xml.write('\t\t<segmented>0</segmented>\n')# 节点判断for multi in json_file["shapes"]:points = np.array(multi["points"])xmin = min(points[:, 0])xmax = max(points[:, 0])ymin = min(points[:, 1])ymax = max(points[:, 1])label = multi["label"]if xmax <= xmin:passelif ymax <= ymin:passelse:xml.write('\t<object>\n')xml.write('\t\t<name>' + json_file["shapes"][i]["label"] + '</name>\n')xml.write('\t\t<pose>Unspecified</pose>\n')xml.write('\t\t<truncated>0</truncated>\n')xml.write('\t\t<difficult>0</difficult>\n')xml.write('\t\t<bndbox>\n')xml.write('\t\t\t<xmin>' + str(xmin) + '</xmin>\n')xml.write('\t\t\t<ymin>' + str(ymin) + '</ymin>\n')xml.write('\t\t\t<xmax>' + str(xmax) + '</xmax>\n')xml.write('\t\t\t<ymax>' + str(ymax) + '</ymax>\n')xml.write('\t\t</bndbox>\n')xml.write('\t</object>\n')print(json_filename, xmin, ymin, xmax, ymax, label)i = i + 1xml.write('</annotation>')# 复制图片newImgName = os.path.join('D:/something/homework/deeplearning/Mydataset/data/images', img_name.split("\\")[-1])shutil.copy(img_name, newImgName)

自行修改地址，还有倒数第二行地址。

结果：xml文件在如下图：

4、划分数据集

(注释是因为一开始用相对路径有问题，后面全部改成绝对路径了）

import os
import randomtrainval_percent = 0.9
train_percent = 0.9
# 1. 获取脚本所在目录作为基础路径
base_dir = os.path.dirname(os.path.abspath(__file__))# 2. 定义所有路径（基于脚本所在目录）
xmlfilepath = os.path.join(base_dir, 'data', 'Annotations')
imagesets_dir = os.path.join(base_dir, 'data', 'ImageSets')# 3. 确保ImageSets目录存在
os.makedirs(imagesets_dir, exist_ok=True)# 4. 列出XML文件（添加错误检查）
try:total_xml = os.listdir(xmlfilepath)
except FileNotFoundError:print(f"错误：目录不存在 - {xmlfilepath}")exit(1)# 5. 避免使用关键字"list"作为变量名
file_count = len(total_xml)
file_indices = range(file_count)# 6. 计算分割点
trainval_percent = 0.9
train_percent = 0.9
tv = int(file_count * trainval_percent)
tr = int(tv * train_percent)# xmlfilepath = 'data/Annotations'
# total_xml = os.listdir(xmlfilepath)# num = len(total_xml)
# list = range(num)
# tv = int(num * trainval_percent)
# tr = int(tv * train_percent)# 7. 创建样本集
trainval = random.sample(file_indices, tv)
train = random.sample(trainval, tr)# 8. 创建文件路径（使用完整路径）
trainval_path = os.path.join(imagesets_dir, 'trainval.txt')
test_path = os.path.join(imagesets_dir, 'test.txt')
train_path = os.path.join(imagesets_dir, 'train.txt')
val_path = os.path.join(imagesets_dir, 'val.txt')# 9. 安全地打开文件（使用with语句自动关闭）
with open(trainval_path, 'w') as ftrainval, \open(test_path, 'w') as ftest, \open(train_path, 'w') as ftrain, \open(val_path, 'w') as fval:for i in file_indices:# 10. 提取文件名（不含扩展名）file_name = os.path.splitext(total_xml[i])[0]if i in trainval:ftrainval.write(file_name + '\n')if i in train:ftrain.write(file_name + '\n')else:fval.write(file_name + '\n')else:ftest.write(file_name + '\n')print("数据集分割完成！")
print(f"总文件数: {file_count}")
print(f"训练验证集: {tv} 个文件")
print(f"测试集: {file_count - tv} 个文件")
print(f"训练集: {tr} 个文件")
print(f"验证集: {tv - tr} 个文件")# trainval = random.sample(list, tv)
# train = random.sample(trainval, tr)# ftrainval = open('data/ImageSets/trainval.txt', 'w')
# ftest = open('data/ImageSets/test.txt', 'w')
# ftrain = open('data/ImageSets/train.txt', 'w')
# fval = open('data/ImageSets/val.txt', 'w')# for i in list:
#     name = total_xml[i][:-4] + '\n'
#     if i in trainval:
#         ftrainval.write(name)
#         if i in train:
#             ftrain.write(name)
#         else:
#             fval.write(name)
#     else:
#         ftest.write(name)# ftrainval.close()
# ftrain.close()
# fval.close()
# ftest.close()

结果：（里面保存的就是划分图片的名字）

5、xml转txt

yolo用得是txt格式的数据，要进行归一化等，如下：

# # xml解析包
# import xml.etree.ElementTree as ET
# import pickle
# import os
# from os import listdir, getcwd
# from os.path import join# sets = ['train', 'test', 'val']
# # classes = ['round', 'square', 'rectangular', 'flat_elliptical', 'mantou', 'groove', 'regular_triangle', 'd_shaped', 'angle']
# classes=['0', '1', '2', '3', '4', '5']  # 类别名称，0-5对应类别# # 进行归一化操作
# def convert(size, box): # size:(原图w,原图h) , box:(xmin,xmax,ymin,ymax)
#     dw = 1./size[0]     # 1/w
#     dh = 1./size[1]     # 1/h
#     x = (box[0] + box[1])/2.0   # 物体在图中的中心点x坐标
#     y = (box[2] + box[3])/2.0   # 物体在图中的中心点y坐标
#     w = box[1] - box[0]         # 物体实际像素宽度
#     h = box[3] - box[2]         # 物体实际像素高度
#     x = x*dw    # 物体中心点x的坐标比(相当于 x/原图w)
#     w = w*dw    # 物体宽度的宽度比(相当于 w/原图w)
#     y = y*dh    # 物体中心点y的坐标比(相当于 y/原图h)
#     h = h*dh    # 物体宽度的宽度比(相当于 h/原图h)
#     return (x, y, w, h)    # 返回 相对于原图的物体中心点的x坐标比,y坐标比,宽度比,高度比,取值范围[0-1]# # year ='2012', 对应图片的id（文件名）
# def convert_annotation(image_id):
#     '''
#     将对应文件名的xml文件转化为label文件，xml文件包含了对应的bunding框以及图片长款大小等信息，
#     通过对其解析，然后进行归一化最终读到label文件中去，也就是说
#     一张图片文件对应一个xml文件，然后通过解析和归一化，能够将对应的信息保存到唯一一个label文件中去
#     labal文件中的格式：calss x y w h　　同时，一张图片对应的类别有多个，所以对应的ｂｕｎｄｉｎｇ的信息也有多个
#     '''
#     # 对应的通过year 找到相应的文件夹，并且打开相应image_id的xml文件，其对应bund文件
#     in_file = open('data/Annotations/%s.xml' % (image_id), encoding='utf-8')
#     # 准备在对应的image_id 中写入对应的label，分别为
#     # <object-class> <x> <y> <width> <height>
#     out_file = open('data/labels/%s.txt' % (image_id), 'w', encoding='utf-8')
#     # 解析xml文件
#     tree = ET.parse(in_file)
#     # 获得对应的键值对
#     root = tree.getroot()
#     # 获得图片的尺寸大小
#     size = root.find('size')
#     # 如果xml内的标记为空，增加判断条件
#     if size != None:
#         # 获得宽
#         w = int(size.find('width').text)
#         # 获得高
#         h = int(size.find('height').text)
#         # 遍历目标obj
#         for obj in root.iter('object'):
#             # 获得difficult ？？
#             difficult = obj.find('difficult').text
#             # 获得类别 =string 类型
#             cls = obj.find('name').text
#             # 如果类别不是对应在我们预定好的class文件中，或difficult==1则跳过
#             if cls not in classes or int(difficult) == 1:
#                 continue
#             # 通过类别名称找到id
#             cls_id = classes.index(cls)
#             # 找到bndbox 对象
#             xmlbox = obj.find('bndbox')
#             # 获取对应的bndbox的数组 = ['xmin','xmax','ymin','ymax']
#             b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
#                  float(xmlbox.find('ymax').text))
#             print(image_id, cls, b)
#             # 带入进行归一化操作
#             # w = 宽, h = 高， b= bndbox的数组 = ['xmin','xmax','ymin','ymax']
#             bb = convert((w, h), b)
#             # bb 对应的是归一化后的(x,y,w,h)
#             # 生成 calss x y w h 在label文件中
#             out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')# # 返回当前工作目录
# wd = getcwd()
# print(wd)# for image_set in sets:
#     '''
#     对所有的文件数据集进行遍历
#     做了两个工作：
# 　　　　１．将所有图片文件都遍历一遍，并且将其所有的全路径都写在对应的txt文件中去，方便定位
# 　　　　２．同时对所有的图片文件进行解析和转化，将其对应的bundingbox 以及类别的信息全部解析写到label 文件中去
#     　　　　　最后再通过直接读取文件，就能找到对应的label 信息
#     '''
#     # 先找labels文件夹如果不存在则创建
#     if not os.path.exists('data/labels/'):
#         os.makedirs('data/labels/')
#     # 读取在ImageSets/Main 中的train、test..等文件的内容
#     # 包含对应的文件名称
#     image_ids = open('data/ImageSets/%s.txt' % (image_set)).read().strip().split()
#     # 打开对应的2012_train.txt 文件对其进行写入准备
#     list_file = open('data/%s.txt' % (image_set), 'w')
#     # 将对应的文件_id以及全路径写进去并换行
#     for image_id in image_ids:
#         list_file.write('E:/Mydataset/data/images/%s.bmp\n' % (image_id))
#         # 调用  year = 年份  image_id = 对应的文件名_id
#         convert_annotation(image_id)
#     # 关闭文件
#     list_file.close()import xml.etree.ElementTree as ET
import pickle
import os
import random
from os import listdir, getcwd
from os.path import join, abspath, dirname# 0. 获取脚本所在目录作为基础路径
base_dir = dirname(abspath(__file__))# 1. 创建所有必要的目录
required_dirs = ['data/Annotations','data/labels','data/ImageSets','data/images'
]for dir_path in required_dirs:full_path = join(base_dir, dir_path)os.makedirs(full_path, exist_ok=True)print(f"已创建目录: {full_path}")# 2. 定义数据集和类别
sets = ['train', 'test', 'val']
classes = ['0', '1', '2', '3', '4', '5']  # 类别名称# 3. 归一化函数（保持不变）
def convert(size, box):dw = 1./size[0]dh = 1./size[1]x = (box[0] + box[1])/2.0y = (box[2] + box[3])/2.0w = box[1] - box[0]h = box[3] - box[2]x = x*dww = w*dwy = y*dhh = h*dhreturn (x, y, w, h)# 4. 修改convert_annotation函数使用绝对路径
def convert_annotation(image_id):# 使用绝对路径xml_path = join(base_dir, 'data', 'Annotations', f'{image_id}.xml')label_path = join(base_dir, 'data', 'labels', f'{image_id}.txt')# 检查XML文件是否存在if not os.path.exists(xml_path):print(f"警告: XML文件不存在 - {xml_path}")returntry:tree = ET.parse(xml_path)root = tree.getroot()size = root.find('size')if size is None:print(f"警告: {xml_path} 中没有找到size元素")returnw = int(size.find('width').text)h = int(size.find('height').text)with open(label_path, 'w', encoding='utf-8') as out_file:for obj in root.iter('object'):difficult = obj.find('difficult').textcls = obj.find('name').textif cls not in classes or (difficult and 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))bb = convert((w, h), b)out_file.write(f"{cls_id} {' '.join(map(str, bb))}\n")print(f"已处理: {image_id}")except Exception as e:print(f"处理 {xml_path} 时出错: {str(e)}")# 5. 主程序
if __name__ == "__main__":print(f"脚本目录: {base_dir}")print(f"当前工作目录: {getcwd()}")# 检查ImageSets目录中的txt文件是否存在imagesets_dir = join(base_dir, 'data', 'ImageSets')for set_name in sets:txt_path = join(imagesets_dir, f'{set_name}.txt')if not os.path.exists(txt_path):print(f"警告: {txt_path} 不存在，将自动创建空的")open(txt_path, 'w').close()# 处理每个数据集for image_set in sets:txt_path = join(imagesets_dir, f'{image_set}.txt')try:with open(txt_path, 'r', encoding='utf-8') as f:image_ids = [line.strip() for line in f.readlines() if line.strip()]if not image_ids:print(f"警告: {txt_path} 是空的，跳过")continueprint(f"处理 {image_set} 集，共 {len(image_ids)} 个图像")list_file_path = join(base_dir, 'data', f'{image_set}.txt')with open(list_file_path, 'w', encoding='utf-8') as list_file:for image_id in image_ids:# 使用相对路径或绝对路径# 相对路径（推荐，便于在不同机器上使用）# image_path = join('data', 'images', f'{image_id}.bmp')# 绝对路径（根据实际情况修改）image_path = join(base_dir, 'data', 'images', f'{image_id}.jpg')list_file.write(f"{image_path}\n")# 处理标注convert_annotation(image_id)print(f"完成 {image_set} 集处理")except Exception as e:print(f"处理 {image_set} 集时出错: {str(e)}")print("所有处理完成！")

结果：

数据集分割完成！

总文件数: 100

训练验证集: 90 个文件

测试集: 10 个文件

训练集: 81 个文件

验证集: 9 个文件

6、数据类别

运行脚本知道甲状腺结节类别：['1', '3', '5', '4', '2', '0']

脚本如下（如果知道自己有哪些类别的，就可以不用运行如下代码）：

import xml.dom.minidom as xmldom
import osbase_dir = os.path.dirname(os.path.abspath(__file__))# 2. 定义所有路径（基于脚本所在目录）
annotation_path = os.path.join(base_dir, 'data', 'Annotations')
# voc数据集获取所有标签的所有类别数"
# annotation_path = "./data/Annotations/"annotation_names = [os.path.join(annotation_path, i) for i in os.listdir(annotation_path)]labels = list()
for names in annotation_names:xmlfilepath = namesdomobj = xmldom.parse(xmlfilepath)# 得到元素对象elementobj = domobj.documentElement# 获得子标签subElementObj = elementobj.getElementsByTagName("object")for s in subElementObj:label = s.getElementsByTagName("name")[0].firstChild.data# print(label)if label not in labels:labels.append(label)
print(labels)