yolov5多个数据集合并

个体学习户

于 2024-10-11 00:52:17 发布

阅读量567

点赞数 16

文章标签： YOLO

本文链接：https://blog.csdn.net/weixin_71704287/article/details/142799680

版权

依旧是学习用途，本期主要内容：

1、数据集文件夹合并

2、数据集标签文件重写

以下是正文

一、将数据集文件夹合并

1、将需要的数据集们准备好，格式如下：

其中data.yaml和train文件夹是必要的，train文件夹内容为：

images文件夹内均为用于训练的图片，labels为txt标签文件。

2、需要合并数据集格式均如上所述（train文件为必须，其中包含图片和标签）

将train文件夹内的images和labels文件夹分别合并，即将内容都复制黏贴到一起，可以额外准备一个文件夹，将你要合并的数据集内容都放里面。（没有valid和test文件也别急，参考我这篇博客，里面有从train文件夹中抽出valid和test文件夹内容的python代码）

例如我这里将①c.v1i.yolov5pytorch数据集和②ship_detec.v1i.yolov5pytorch数据集都放入③dataset_test数据集中。如图所示：

其中，data.yaml文件先不用管，后续会修改。

检查一下文件夹基础内容是否合并成功，从图中可以看到已经合并成功了，即两种图片都有。

二、修改数据集的标签内容

1、随便打开某个数据集的某个标签，例如我这里：

"D:\dataset_test\train\labels\0_jpeg.rf.5baa289e18c78cde890157c41f4e2b76.txt"，我们可以看到：

其中，最左边的数字代表种类，后续的数字代表别的信息，我们这里可以不用管，只需要对最左边的数字操作即可。

2、假如只有两个数据集，那么一个数据集txt标签的数字可以不改变，只修改另一个数据集复制过来的txt标签文件的数字即可。

用我这个数据集举例，我这里两个数据集分别是军舰和雷达的，并且都只有一种种类。

可以看到其中193_jpeg类似的文件属于雷达的labels，而tau1类似文件属于军舰的labels。

我们这边可以保留一种labels，修改另一种labels。假设雷达的种类是0，已经在上面的txt文件展示过了，那么军舰就不能是0了，我们需要给它们改成1，如图所示。但是我们不可能一个一个文件打开去修改，因此就需要python脚本辅助了。

3、编写修改脚本，我这里参考了这个博客的代码

但是做了一些修改，代码如下：

import os

YOLOV5_LABEL_ROOT = r"C:\Users\19216\Desktop\army-vehicle.v1-truck-tank-armoured-artillery-person.yolov5pytorch\train\labels\\"  # yolov5 导出的推理图片的 txt
DATASET_LABEL_ROOT = r"D:\dataset_test\train\labels\\"  # 数据集的路径

# 定义替换规则
replacement_rules = {
    '0': '1'
}

if __name__ == '__main__':
    yolo_files = os.listdir(YOLOV5_LABEL_ROOT)

    # 遍历文件里面有 .txt 结尾的
    for file_name in yolo_files:
        # 判断 txt 文件才进行读取
        if not file_name.endswith(".txt"):
            continue

        file_path = os.path.join(YOLOV5_LABEL_ROOT, file_name)  # 使用 os.path.join 拼接路径
        try:
            with open(file_path, "r") as f:
                lines = f.readlines()
                modified_lines = []

                # 根据替换规则修改每一行的第一个字符
                for line in lines:
                    first_char = line[0]
                    if first_char in replacement_rules:
                        new_first_char = replacement_rules[first_char]
                        modified_line = new_first_char + line[1:]
                        modified_lines.append(modified_line)
                    else:
                        # 如果第一个字符不在替换规则中，保留原样
                        modified_lines.append(line)

                data_path = os.path.join(DATASET_LABEL_ROOT, file_name)
                print(data_path)

                # 清空目标文件（如果已存在）
                with open(data_path, "w") as fd:
                    pass

                # 将修改后的内容写入数据集的标注文件
                with open(data_path, "w") as fd:
                    for modified_line in modified_lines:
                        fd.write(modified_line)
        except Exception as e:
            print(f"Error processing {file_path}: {str(e)}")

这个代码的作用就是，把原数据集labels的txt文件的内容，修改后根据相同名字文件进行覆盖进目标文件夹。可以看到我这里是把原文件中每行首字符为0的都置1，从而达到区分的目的。