【DL】（数据）数据集的划分

【DL】（数据）数据集的划分

一、划分形式

数据集的划分工作是在模型训练之前就要完成的，我们得到的数据集大多都只放在一个文件夹中，标签存放在一个csv文件中，数据集的划分需要我们自己来完成，我们的目标是将一个文件夹中的所有图片，按照比例划分为训练集和验证集，在训练集和验证集中，每一个子文件夹的名称表示该子文件夹中图片的标签。
以下图二分类为例：（0表示负类，1表示正类）

二、思路：

1、读取csv文件，将里面所有图片的文件名存储在列表filename中，每个文件名对应的标签存储在列表label中。
2、先创建train, valid文件夹
3、对于train和valid，创建其子文件夹用来表示类别
4、对filename和label进行切分，得到表示训练集的train_img, train_lab和表示验证集的valid_img, valid_lab。
5、遍历train_img和valid_img, 对当前图片的标签进行判断，若为0，则添加到对应的train/valid中的0/1文件夹中。

三、代码

import os
import shutil
import pandas as pd
import numpy as np

# ====对数据的csv文件进行处理，将图片的文件名存入face_filename,标签存入face_label =======================================
data = pd.read_csv("D:\\Pytorch\\pythonProject\\Pycharm_scripts_torch\\xxx\\train.labels.csv")
data = data.iloc[:, 0]
data = np.array(data)
# print(data)
face_filename = []
face_label = []
for i in range(len(data)):
    face_filename.append(data[i].split("\t")[0])
    face_label.append(data[i].split("\t")[1])
# print(face_filename)
# print(face_label)
# ===================================================================================================================

base_dir = os.path.abspath(r"D:\Pytorch\pythonProject\Pycharm_scripts_torch\xxx\image\trainvalid")

# os.path.makedir  用来创建新的文件夹
def makedir(new_dir):
    if not os.path.exists(new_dir):
        os.makedirs(new_dir)

# 创建train、valid文件夹
makedir(r"D:\Pytorch\pythonProject\Pycharm_scripts_torch\xxx\image\train")
makedir(r"D:\Pytorch\pythonProject\Pycharm_scripts_torch\xxx\image\valid")
train_path = os.path.abspath(r"D:\Pytorch\pythonProject\Pycharm_scripts_torch\xxx\image\train")
valid_path = os.path.abspath(r"D:\Pytorch\pythonProject\Pycharm_scripts_torch\xxx\image\valid")

# 划分比例
train_percent = 0.9
valid_percent = 0.1

train_num = len(face_filename) * train_percent
valid_num = len(face_filename) * valid_percent

train_img = face_filename[:int(train_num)]            # 列表，存储训练集图片的文件名
train_lab = face_label[:int(train_num)]               # 列表，存储训练集图片的标签
valid_img = face_filename[int(train_num):]            # 列表，存储验证集图片的文件名
valid_lab = face_label[int(train_num):]               # 列表，存储验证集图片的标签

# 划分训练集
makedir(os.path.join(train_path, str(0)))    # 创建正类文件
makedir(os.path.join(train_path, str(1)))    # 创建负类文件

for i in range(len(train_img)):
    img_path = os.path.join(base_dir, train_img[i])
    img_label = train_lab[i]
    if img_label == '0':
        shutil.copy(img_path, os.path.join(train_path, str(0)))  # shutil.copy(a,b) 将a文件复制到b文件中，具体参数请查阅资料
    elif img_label == '1':
        shutil.copy(img_path, os.path.join(train_path, str(1)))

# 划分验证集
makedir(os.path.join(valid_path, str(0)))
makedir(os.path.join(valid_path, str(1)))

for i in range(len(valid_img)):
    img_path = os.path.join(base_dir, valid_img[i])
    img_label = valid_lab[i]
    if img_label == '0':
        shutil.copy(img_path, os.path.join(valid_path, str(0)))  
    elif img_label == '1':
        shutil.copy(img_path, os.path.join(valid_path, str(1)))