（1）数据集的划分

数据集划分与PNG转JPEG

原创已于 2022-10-11 10:23:52 修改 · 452 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#python #深度学习 #机器学习

于 2022-10-11 09:14:15 首次发布

torch 专栏收录该内容

32 篇文章

订阅专栏

本文介绍了如何按比例划分数据集为train、val和test，并提供了基础转换代码，特别是详细讲解了将PNG图像转换为JPEG格式的过程，适用于深度学习和机器学习的数据预处理。

1、数据集的划分

将数据集根据需要，按比例划分为train，val和test

生成后的文件夹，可以采用pytorch使用 torchvision.datasets.folder.ImageFolder(root='data/train')

直接读取

#--data_floder_orign
#  --class1
#     --001.jpg
#     --002.jpg  
#     --003.jpg  
#  --class2
#     --004.jpg
#     --005.jpg  
#     --006.jpg  
#  --class3
#     --007.jpg
#     --008.jpg  
#     --009.jpg 
#---------------------------------------------
#根据比例产生
# -data
#    --train
#      --class1
#         --001.jpg
#      --class2
#         --006.jpg 
#      --class3
#         --008.jpg
#    --val
#         --002.jpg 
#      --class2
#         --005.jpg
#      --class3
#         --009.jpg 
#    --test
#      --class1
#         --003.jpg
#      --class2
#         --004.jpg
#      --class3
#         --007.jpg
#######################################

2、基础转换代码的实现

# *_*coding: utf-8 *_*
#######################################
#函数功能
#对照原始文件，按比例复制生成train,val,test文件夹

#--data_floder_orign
#  --class1
#     --001.jpg
#     --002.jpg  
#     --003.jpg  
#  --class2
#     --004.jpg
#     --005.jpg  
#     --006.jpg  
#  --class3
#     --007.jpg
#     --008.jpg  
#     --009.jpg 
#---------------------------------------------
#根据比例产生
# -data
#    --train
#      --class1
#         --001.jpg
#      --class2
#         --006.jpg 
#      --class3
#         --008.jpg
#    --val
#         --002.jpg 
#      --class2
#         --005.jpg
#      --class3
#         --009.jpg 
#    --test
#      --class1
#         --003.jpg
#      --class2
#         --004.jpg
#      --class3
#         --007.jpg
#######################################


import os
import random
import shutil  #复制文件用到
import time

"""
    函数说明:获取指定文件夹下的目录文件的函数，陈20221011
    :param    dir_path: 目录的路径
    :return:  dir_list: 该目录下的文件夹名称，list类型
 """
def  get_dir(dir_path):
    dir_list = []
    for filename in sorted(os.listdir(dir_path)):
        #只判断目录的
        if os.path.isdir(os.path.join(dir_path,filename)):
            dir_list.append(filename)
    return dir_list

"""
    函数说明: 复制文件的函数
    :param    src_file: 原始文件
    :return:  dst_path: 目标目录
 """
def copy_file(src_file,dst_dir_path):                       # 复制函数
    if not os.path.isfile(src_file):
        print ("%s not exist!"%(src_file))
    else:
        fpath,fname=os.path.split(src_file)             # 分离文件名和路径
        if not os.path.exists(dst_dir_path):
            os.makedirs(dst_dir_path)                       # 创建路径
        shutil.copy(src_file, dst_dir_path + fname)          # 复制文件
        print ("copy %s -> %s"%(src_file, dst_dir_path + fname))


"""
    函数说明: 将指定文件列表中的所有文件拷贝到指定的文件夹下如果文件夹
             不存在则自动创建，陈20221011
    :param    dest_file_dir: 目标目录的路径
    :param    origin_file_dir: 原始待拷贝目录的路径
    :param    origion_file_list: 原始待拷贝文件的文件名，注意该列表不包含文件的路径
    :return:  
 """
def copy_File_To_Dest_directory(dest_file_dir, origin_file_dir, origion_file_list):
    # print("dest_file_dir: ",dest_file_dir)
    for file_name in iter(origion_file_list):
        file_name_full_path = origin_file_dir + file_name
        # print("chen copy file")
        # print(file_name_full_path)
        copy_file(file_name_full_path,dest_file_dir)

"""
    函数说明: 对指定的分类按比例进行分割，并复制到指定文件夹，陈20221011
    :param    origion_path: 原始目录的路径
    :param    save_path: 保存目录的路径
    :param    calss_name: 待处理分类的名称
    :param    train_rate: 训练数据比例
    :param    val_rate: 待校验数据比例
    :param    test_rate: 待测试数据比例
    :return:  
"""
def dataset_Classification(origion_path,save_path,calss_name,train_rate,val_rate,test_rate):
    #（1）首先获取该分类目录下的所有文件
    class_path = origion_path + calss_name + "/"
    all_class_file_name = []
    for file_name in sorted(os.listdir(os.path.join(class_path))):
        #考虑全面起见，过滤掉文件夹，只剩下文件
        if os.path.isdir(os.path.join(class_path, file_name)):
            continue
        all_class_file_name.append(file_name)

    print("原始的数据分布: ")
    print(all_class_file_name)

    #（2）对获取的分类文件list进行置乱
    random.shuffle(all_class_file_name)
    print("置乱后的数据分别: ")
    print(all_class_file_name)

    #（3）按比例设置train_file_list, val_file_list, test_file_list
    total_file_num = len(all_class_file_name)
    print("total: ",calss_name, " 个数 ",total_file_num)

    train_file_list = all_class_file_name[:int(train_rate*total_file_num)]
    # print("train_file_list: ",len(train_file_list),train_file_list)

    val_file_list = all_class_file_name[int(train_rate*total_file_num):int((train_rate+val_rate)*total_file_num)]
    # print("val_file_list: ",len(val_file_list),val_file_list)

    test_file_list = all_class_file_name[int((train_rate+val_rate)*total_file_num):]
    # print("val_file_list: ",len(test_file_list),test_file_list)

    #(4)复制到指定的目录
    copy_File_To_Dest_directory(save_path+"train/"+calss_name+"/",class_path,train_file_list)
    copy_File_To_Dest_directory(save_path+"val/"+calss_name+"/",class_path,val_file_list)
    copy_File_To_Dest_directory(save_path+"test/"+calss_name+"/",class_path,test_file_list)



if __name__ == '__main__':
    time_start = time.time()

    # 原始按文件夹分布的数据集路径
    # origion_path = './pokeman/'
    origion_path = 'D:/pytorch_learning2022/2chensample_code/pokeman/'
    # 调整后按文件夹分布的数据集路径
    # save_path = './data'
    save_path = 'D:/pytorch_learning2022/2chensample_code/data/'

    # 训练集、验证集、测试集的比例
    train_rate = 0.6
    val_rate = 0.2
    test_rate = 0.2

    # 数据集类别及数量
    all_class_list = get_dir(origion_path)
    print("共有分类个数： ",len(all_class_list))
    print(all_class_list)

    # dataset_Classification(origion_path,save_path,all_class_list[0],train_rate,val_rate,test_rate)
    #遍历所有分类
    print('------开始拷贝---------')
    for calss_name in iter(all_class_list):
        print(calss_name,"拷贝中")
        dataset_Classification(origion_path,save_path,calss_name,train_rate,val_rate,test_rate)
        

    print('划分完毕!')
    time_end = time.time()
    print('---------------')
    print('训练集和测试集划分共耗时%s!' % (time_end - time_start))

3、将png转换为jpg的升级代码

网络数据集的很多照片是png，需要转换为jpg，更正规

3.1 png的情况

        #注意两种情况
        #（1）部分文件后缀名是jpg或者jpeg，实际上是png
        #（2）注意部分文件实际是jpg或者jpeg，但是文件名命名为.png，强迫症，也转换一下

3.2 包括了png转jpeg的代码

# *_*coding: utf-8 *_*
#######################################
#函数功能
#对照原始文件，按比例复制生成train,val,test文件夹
#同时把png文件转换为jpg


#--data_floder_orign
#  --class1
#     --001.jpg
#     --002.jpg  
#     --003.jpg  
#  --class2
#     --004.jpg
#     --005.jpg  
#     --006.jpg  
#  --class3
#     --007.jpg
#     --008.jpg  
#     --009.jpg 
#---------------------------------------------
#根据比例产生
# -data
#    --train
#      --class1
#         --001.jpg
#      --class2
#         --006.jpg 
#      --class3
#         --008.jpg
#    --val
#         --002.jpg 
#      --class2
#         --005.jpg
#      --class3
#         --009.jpg 
#    --test
#      --class1
#         --003.jpg
#      --class2
#         --004.jpg
#      --class3
#         --007.jpg
#######################################


import os
import random
import shutil  #复制文件用到
import time
import imghdr
from    PIL import Image

"""
    函数说明:获取指定文件夹下的目录文件的函数，陈20221011
    :param    dir_path: 目录的路径
    :return:  dir_list: 该目录下的文件夹名称，list类型
 """
def  get_dir(dir_path):
    dir_list = []
    for filename in sorted(os.listdir(dir_path)):
        #只判断目录的
        if os.path.isdir(os.path.join(dir_path,filename)):
            dir_list.append(filename)
    return dir_list

"""
    函数说明: 复制文件的函数
    :param    src_file: 原始文件
    :return:  dst_path: 目标目录
 """
def copy_file(src_file,dst_dir_path):                       # 复制函数
    if not os.path.isfile(src_file):
        print ("%s not exist!"%(src_file))
    else:
        fpath,fname=os.path.split(src_file)             # 分离文件名和路径
        if not os.path.exists(dst_dir_path):
            os.makedirs(dst_dir_path)                       # 创建路径
        shutil.copy(src_file, dst_dir_path + fname)          # 复制文件
        print ("copy %s -> %s"%(src_file, dst_dir_path + fname))

"""
    函数说明: 复制文件的函数，注意：同时判断是否是png，如果是的话转换为jpg
    :param    src_file: 原始文件
    :return:  dst_path: 目标目录
 """
def copy_file_judge_png_to_jpg(src_file,dst_dir_path):                       # 复制函数
    if not os.path.isfile(src_file):
        print ("%s not exist!"%(src_file))
    else:
        imgType = imghdr.what(src_file)
        print(imgType)
        #注意两种情况
        #（1）部分文件后缀名是jpg或者jpeg，实际上是png
        #（2）注意部分文件实际是jpg或者jpeg，但是文件名命名为.png，强迫症，也转换一下
        suffix_only = os.path.splitext(src_file)[-1]

        if not os.path.exists(dst_dir_path):
            os.makedirs(dst_dir_path)                       # 创建路径
        #如果是png转换后存储
        if imgType == 'png' or suffix_only=='.png':
            print("this is png")
            im = Image.open(src_file)
            rgb_im = im.convert('RGB')
            file_suffix = src_file.split(".png")[0] +".jpg"
            fpath,fname=os.path.split(file_suffix)             # 分离文件名和路径
            rgb_im.save(dst_dir_path + fname)
            print ("change png to jpg %s -> %s"%(src_file, dst_dir_path + fname))
        else:
            fpath,fname=os.path.split(src_file)             # 分离文件名和路径
            # if not os.path.exists(dst_dir_path):
            #     os.makedirs(dst_dir_path)                       # 创建路径
            shutil.copy(src_file, dst_dir_path + fname)          # 复制文件
            print ("copy %s -> %s"%(src_file, dst_dir_path + fname))


"""
    函数说明: 将指定文件列表中的所有文件拷贝到指定的文件夹下如果文件夹
             不存在则自动创建，陈20221011
    :param    dest_file_dir: 目标目录的路径
    :param    origin_file_dir: 原始待拷贝目录的路径
    :param    origion_file_list: 原始待拷贝文件的文件名，注意该列表不包含文件的路径
    :return:  
 """
def copy_File_To_Dest_directory(dest_file_dir, origin_file_dir, origion_file_list):
    # print("dest_file_dir: ",dest_file_dir)
    for file_name in iter(origion_file_list):
        file_name_full_path = origin_file_dir + file_name
        # print("chen copy file")
        # print(file_name_full_path)
        # copy_file(file_name_full_path,dest_file_dir)
        copy_file_judge_png_to_jpg(file_name_full_path,dest_file_dir)
        

"""
    函数说明: 对指定的分类按比例进行分割，并复制到指定文件夹，陈20221011
    :param    origion_path: 原始目录的路径
    :param    save_path: 保存目录的路径
    :param    calss_name: 待处理分类的名称
    :param    train_rate: 训练数据比例
    :param    val_rate: 待校验数据比例
    :param    test_rate: 待测试数据比例
    :return:  
"""
def dataset_Classification(origion_path,save_path,calss_name,train_rate,val_rate,test_rate):
    #（1）首先获取该分类目录下的所有文件
    class_path = origion_path + calss_name + "/"
    all_class_file_name = []
    for file_name in sorted(os.listdir(os.path.join(class_path))):
        #考虑全面起见，过滤掉文件夹，只剩下文件
        if os.path.isdir(os.path.join(class_path, file_name)):
            continue
        all_class_file_name.append(file_name)

    print("原始的数据分布: ")
    print(all_class_file_name)

    #（2）对获取的分类文件list进行置乱
    random.shuffle(all_class_file_name)
    print("置乱后的数据分别: ")
    print(all_class_file_name)

    #（3）按比例设置train_file_list, val_file_list, test_file_list
    total_file_num = len(all_class_file_name)
    print("total: ",calss_name, " 个数 ",total_file_num)

    train_file_list = all_class_file_name[:int(train_rate*total_file_num)]
    # print("train_file_list: ",len(train_file_list),train_file_list)

    val_file_list = all_class_file_name[int(train_rate*total_file_num):int((train_rate+val_rate)*total_file_num)]
    # print("val_file_list: ",len(val_file_list),val_file_list)

    test_file_list = all_class_file_name[int((train_rate+val_rate)*total_file_num):]
    # print("val_file_list: ",len(test_file_list),test_file_list)

    #(4)复制到指定的目录
    copy_File_To_Dest_directory(save_path+"train/"+calss_name+"/",class_path,train_file_list)
    copy_File_To_Dest_directory(save_path+"val/"+calss_name+"/",class_path,val_file_list)
    copy_File_To_Dest_directory(save_path+"test/"+calss_name+"/",class_path,test_file_list)

def test_fun22():
    image_file = "D:/pytorch_learning2022/data/pokeman/train/mewtwo/00000008.png"
    imgType = imghdr.what(image_file)
    print(imgType)

    suffix_only = os.path.splitext(image_file)[-1]
    print(suffix_only)

    if suffix_only == '.png':
        print("this is png")

def test_fun():
    #判断是否为png图
    image_file = "D:/pytorch_learning2022/data/pokeman/train/bulbasaur/00000002.png"
    imgType = imghdr.what(image_file)
    print(imgType)
    if imgType == 'png':
        print("this is png")
    im = Image.open(image_file)
    rgb_im = im.convert('RGB')
    # rgb_im.save()

    suffix = image_file.split(".png")[0]
    print(suffix)
    image_file22 = "D:/pytorch_learning2022/data/pokeman/train/bulbasaur/0.0000002.png"
    suffix2 = image_file22.split(".png")[0]
    print(suffix2)


if __name__ == '__main__':

    # test_fun()
    # test_fun22()

    time_start = time.time()

    # 原始按文件夹分布的数据集路径
    # origion_path = './pokeman/'
    origion_path = 'D:/pytorch_learning2022/2chensample_code/pokeman/'
    # 调整后按文件夹分布的数据集路径
    # save_path = './data'
    save_path = 'D:/pytorch_learning2022/data/pokeman/'

    # 训练集、验证集、测试集的比例
    train_rate = 0.6
    val_rate = 0.2
    test_rate = 0.2

    # 数据集类别及数量
    all_class_list = get_dir(origion_path)
    print("共有分类个数： ",len(all_class_list))
    print(all_class_list)

    # dataset_Classification(origion_path,save_path,all_class_list[0],train_rate,val_rate,test_rate)
    #遍历所有分类
    print('------开始拷贝---------')
    for calss_name in iter(all_class_list):
        print(calss_name,"拷贝中")
        dataset_Classification(origion_path,save_path,calss_name,train_rate,val_rate,test_rate)
        

    print('划分完毕!')
    time_end = time.time()
    print('---------------')
    print('训练集和测试集划分共耗时%s!' % (time_end - time_start))