生成分类数据集并划分测试集到EXECL

首先我们要有划分好类别的数据集。

格式如：data---001---类下所有图片

                     ---002---类下所有图片

                     ---003---类下所有图片

                     ---...

已经分好类别后，我们只需要对他7：3划分训练集和测试集，并记录在execl中即可。

很简单其实主要就是读取所有文件夹下文件并划分记录到excel，下面以数据集中一类做例子，之后修改地址和名称即可继续导入其他类别到excel。

#!/usr/bin/python3
# -*- coding: utf-8 -*-

# 导入CSV安装包
import csv
import os
# 1. 创建文件对象
Image_Path=r"F:\ShipRS_sepation_v1\Image"     #图片路径
Annotation_Path=r"F:\ShipRS_sepation_v1\Annotation/"
if not os.path.isdir(Annotation_Path):
    os.makedirs(Annotation_Path)
train = open(Annotation_Path+'/train.csv','w',encoding='utf-8',newline='')
train_writer = csv.writer(train)
train_writer.writerow(["class","label"])
test = open(Annotation_Path+'/test.csv','w',encoding='utf-8',newline='')
test_writer = csv.writer(test)
test_writer.writerow(["class","label"])
count = 0
i = 0
train_num = 0
test_num = 0
imagelist=os.listdir(Image_Path+'/补给舰/')
for image in imagelist:
    (filepath, image_name) = os.path.split(image)
    if i>6:
        test_writer.writerow([image_name, "Supply ship"])
        test_num += 1
    else:
        train_writer.writerow([image_name, "Supply ship"])
        train_num += 1
    i += 1
    if i == 10:
        i = 0
    count += 1
i = 0
print("补给舰训练集 ",train_num)
print("补给舰测试集 ",test_num,'\n')

结果如图