【深度学习】PyTorch如何构建图像数据的测试集？（附代码）

如何构建图像数据的测试集？假设图片已经按照类别分好文件夹，从中间取出一定比例的图片做测试，图片类别标签保存在csv文件。

需求

在每个子文件夹中选取一定数量的图片，移动到test文件夹中，并在test文件夹中进行乱序并重新编号；

需要记录答案，即test中图片的原始对应关系，以的格式，保存在csv中。

代码构建

importosimportmathimportrandomimportcsv
ORI_DIR = './flower_data_new/'TEST_DIR = 'test'TEST_RATE = 0.4SUFFIXES = '.jpg'HEADERS = ['id', 'category']# 计算需要移动到测试集中的图像数量categoryNumList = []categoryNames = os.listdir(ORI_DIR)forcategoryName in categoryNames:categoryDir = os.path.join(ORI_DIR, categoryName)categoryFiles = [filefor file in os.listdir(categoryDir) if file.endswith(SUFFIXES)]categoryNum = len(categoryFiles)categoryNumList.append(categoryNum)selectedNumList = [math.floor(num * TEST_RATE) for num in categoryNumList]testTotalNum = sum(selectedNumList)# 构造乱序数组randomList = [i for i in range(testTotalNum)]fori in range(testTotalNum):rand = random.randint(0, testTotalNum - 1)randomList[i],randomList[rand] = randomList[rand], randomList[i]# 创建测试集文件夹ifnot os.path.exists(TEST_DIR):os.mkdir(TEST_DIR)# 取每类的最后部分，移动到test文件夹，并记录每张图片的id和类别index = 0rows = []for(i, num) in enumerate(selectedNumList):categoryName = categoryNames[i]destDirPath = os.path.join(ORI_DIR, categoryName)files = [filefor file in os.listdir(destDirPath) if file.endswith(SUFFIXES)]forname in files[len(files) - num:]:oriPath = os.path.join(destDirPath, name)destName = str(randomList[index]) + SUFFIXESdestPath = os.path.join(TEST_DIR, destName)rows.append([randomList[index],i])os.rename(oriPath,destPath)index+= 1# 先将rows排序，再存入csvrows = sorted(rows, key=lambda x: x[0])withopen('answer.csv', 'w', newline='') as f:f_csv = csv.writer(f)f_csv.writerow(HEADERS)f_csv.writerows(rows)

此代码假设您的训练文件夹名为 train，测试文件夹名为 test。代码将使用 ImageFolder 获取所有图像路径和对应的标签，并计算需要移动到测试集中的图像数量。然后，它会随机选择需要移动到测试集中的图像，并将它们复制到测试集文件夹中。

在最后一步中，代码将用CSV文件记录每个图像的文件名和标签。可以像使用 answer.csv 文件记录文件名和标签一样将其更改为您想要的格式。

请注意，此代码中使用了 random.sample 方法来随机选择要移动到测试集中的图像。它确保了选中的图像是独立且随机的。