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

发布于:2025-07-10 ⋅ 阅读:(28) ⋅ 点赞:(0)

数据预处理

使用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.dom
 
import 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:
                pass
            elif ymax <= ymin:
                pass
            else:
                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 + 1
        xml.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 random
 
trainval_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  同时,一张图片对应的类别有多个,所以对应的bunding的信息也有多个
#     '''
#     # 对应的通过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:
#     '''
#     对所有的文件数据集进行遍历
#     做了两个工作:
#     1.将所有图片文件都遍历一遍,并且将其所有的全路径都写在对应的txt文件中去,方便定位
#     2.同时对所有的图片文件进行解析和转化,将其对应的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.0
    y = (box[2] + box[3])/2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (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}")
        return
    
    try:
        tree = ET.parse(xml_path)
        root = tree.getroot()
        size = root.find('size')
        
        if size is None:
            print(f"警告: {xml_path} 中没有找到size元素")
            return
            
        w = 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').text
                cls = obj.find('name').text
                
                if cls not in classes or (difficult and int(difficult) == 1):
                    continue
                    
                cls_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} 是空的,跳过")
                continue
                
            print(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 os
 
base_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 = names
    domobj = 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)

开始训练

1、测试

接前文【长文】深度学习小白第一次完整跑通项目全过程,可以先运行一下detect.py进行测试,防止依赖版本等错误

2、参数设置

parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
改为0线程,后面避免问题。


3、写配置文件

(注意冒号后面要有空格)

分别是自己的训练测试验证集地址和类别。

4、加载配置文件

放到代码的data目录下面

在train.py文件中将如下地方改成自己的配置数据集名字。

运行结果:


网站公告

今日签到

点亮在社区的每一天
去签到