python - 将 JPG 和 XML 文件的数据集拆分为训练集和测试集

Question

我有一个用于对象检测算法的数据集，其中包含图片 (.jpg) 和包含边界框的相应 .xml 文件。

我想编写一个脚本，将数据集随机拆分为训练集和测试集，这意味着我必须确保将 jpg 及其相应的 XML 分配到同一目录。

我应该如何编辑以下代码才能实现此目的？

此外，这是执行此操作的“最佳”方法，还是在 xml 到 csv 转换后或生成 csv 到 tfrecords 转换后分割数据集更好？

import shutil, os, glob, random

# List all files in a directory using os.listdir
basepath = '/home/bis/hans/bis/workspace/images/Synced_dataset'
filenames = []

for entry in os.listdir(basepath):
    if os.path.isfile(os.path.join(basepath, entry)):
        #print(entry)
        filenames.append(entry)

filenames.sort()  # make sure that the filenames have a fixed order before shuffling
random.seed(230)
random.shuffle(filenames) # shuffles the ordering of filenames (deterministic given the chosen seed)

split = int(0.8 * len(filenames))
train_filenames = filenames[:split]
test_filenames = filenames[split:]

Answer 1

对我来说最好的选择是创建两个文件列表(filenames for jpg 和 xmlnames for xml >) 以正确的顺序和一个索引列表 indices=[i for i in range(len(filenames))]。

然后您可以随机排列索引列表:

random.seed(230)
random.shuffle(indices)

最后，为 jpg 和 xml 文件创建训练集和测试集:

split = int(0.8 * len(filenames))
file_train = [filenames[idx] for idx in indices[:split]]
file_test = [filenames[idx] for idx in indices[split:]]
xml_train = [xmlnames[idx] for idx in indices[:split]]
xml_test = [xmlnames[idx] for idx in indices[split:]]