把图片数据从文件夹整理成csv文件，并划分训练集、验证集、测试集

代码： 

# -*- encoding: utf-8 -*-
'''
@description： 把图片数据从文件夹整理成csv文件，每一行代表其路径
'''

import pandas as pd 
import os
from PIL import Image

class image2csv(object):
    # 分割训练集 验证集 测试集
    # 做成对应的txt
    def __init__(self,data_root,image_dir,label_dir,slice_data,width_input,height_input):
        self.data_root = data_root
        self.image_dir = image_dir
        self.label_dir = label_dir
        self.slice_train = slice_data[0]
        self.slice_val = slice_data[1]

        self.width = width_input
        self.height = height_input

    def read_path(self):
        images = []
        labels = []
        for i,im in enumerate(os.listdir(self.image_dir)):
            label_name = im.split('.')[0] + '.png'   # 读取图片的名字，去label里面找，确保两个文件夹都有这个名字的图
            if os.path.exists(os.path.join(self.label_dir,label_name)):
                size_w,size_h = Image.open(os.path.join(self.image_dir,im)).size
                size_lw,size_lh = Image.open(os.path.join(self.label_dir,label_name)).size

                if min(size_w,size_lw) >= self.width and min(size_h,size_lh) >= self.height:
                    images.append(os.path.join(self.image_dir,im))
                    labels.append(os.path.join(self.label_dir,label_name))
                else:
                    continue

        assert(len(images)==len(labels)) #

        self.data_length = len(images) # 真正两个文件夹都有的图片的长度

        data_path = {
            'image':images,
            'label':labels,
        }
        return data_path

    def generate_csv(self):
        data_path = self.read_path() # 存放了路径

        data_path_pd = pd.DataFrame(data_path)
        train_slice_point = int(self.slice_train*self.data_length) # 0.7*len
        validation_slice_point = int((self.slice_train+self.slice_val)*self.data_length) # 0.8*len

        train_csv = data_path_pd.iloc[:train_slice_point,:]
        validation_csv = data_path_pd.iloc[train_slice_point:validation_slice_point,:]
        test_csv = data_path_pd.iloc[validation_slice_point:,:]

        train_csv.to_csv(os.path.join(self.data_root,'train.csv'),header=None,index=None)
        validation_csv.to_csv(os.path.join(self.data_root,'val.csv'),header = None,index = None)
        test_csv.to_csv(os.path.join(self.data_root,'test.csv'),header=False,index = False)

if __name__ == "__main__":
    DATA_ROOT =  r'D:\Desktop\segmentation\data'
    image = os.path.join(DATA_ROOT, 'images')
    label = os.path.join(DATA_ROOT, 'labels')
    slice_data = [0.6,0.2,0.2] #  训练 验证 测试所占百分比
    WIDTH = 512
    HEIGHT  =  512
    tocsv = image2csv(DATA_ROOT,image,label,slice_data,WIDTH,HEIGHT)
    tocsv.generate_csv()

 路径截图：