将 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/createview/Vegard/createview/lice_detection_v2/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 用于jpg 和xmlnames 用于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:]]

Answer 2

import shutil, os, glob, random

# List all files in a directory using os.listdir
basepath = 'images/'
labelpath='label/'
filenames = []
xmlnames = []

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

indices=[i for i in range(len(filenames))]        
filenames.sort()
xmlnames.sort() # make sure that the filenames have a fixed order before shuffling
random.seed(230)
random.shuffle(indices) # shuffles the ordering of filenames (deterministic given the chosen seed)

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:]]

print(file_test)
print(xml_test)

所以我按照上述建议（由 Joseph 提出）添加索引，然后当我们制作测试和训练变量时，在变量中添加完全相同的图像和标签，希望这会有所帮助