＜鲁不动AI了＞1.构建自己的图像分类数据集

文章目录

• • 概要
  • 代码讲解
  • 小结

 

概要

图像分类是图像处理的第一步，必须要有数据集，有了数据集，才能分类，构建，训练模型代码。

代码讲解

1.导入必要的库

import os
import shutil
import random
import pandas as pd

        这里用到的库有 os、shutil、random、pandas。其中os用于python库和操作系统进行交互，shutil用于文件的复制、移动等操作，random是用来生成随机数的，pandas用来分析和处理数据。都是必要的库。

2.指定数据集路径并处理

# 指定数据集路径
dataset_path = r"D:\桌面\昆虫"

dataset_name = dataset_path.split('_')[0]
print('数据集', dataset_name)

classes = os.listdir(dataset_path)

len(classes)

classes

        首先读取电脑上存放想训练模型的文件夹，复制路劲后就行了，注意要么在引号前加一个“r”，要么文件路径中使用“\\”双斜杠，否则会报错。

        读取数据集后，我们可以打印出显示一下是否读取成功，如下图。

        同时也可以读取文件夹中所有不同种类的文件夹名称。

ps：没有图片的可以看我主页，我会出一篇图片爬虫的python代码，方便，简单。

3.数据集分类成train和test

        运行代码后，文件夹中会生成两个新文件，一个是train，一个是test。通过设置测试集比例，可以控制训练和测试集的比例。例如这里是0.2，训练集就有4张照片，测试集只有1张。然后的代码就是讲所有种类按比例分配到train和test文件中，形成我们自己的数据集，方便我们后面的训练、制作模型。

# 创建 train 文件夹
os.mkdir(os.path.join(dataset_path, 'train'))

# 创建 test 文件夹
os.mkdir(os.path.join(dataset_path, 'val'))

# 在 train 和 test 文件夹中创建各类别子文件夹
for fruit in classes:
    os.mkdir(os.path.join(dataset_path, 'train', fruit))
    os.mkdir(os.path.join(dataset_path, 'val', fruit))

test_frac = 0.2  # 测试集比例
random.seed(123) # 随机数种子，便于复现

df = pd.DataFrame()

print('{:^18} {:^18} {:^18}'.format('类别', '训练集数据个数', '测试集数据个数'))

for fruit in classes: # 遍历每个类别

    # 读取该类别的所有图像文件名
    old_dir = os.path.join(dataset_path, fruit)
    images_filename = os.listdir(old_dir)
    random.shuffle(images_filename) # 随机打乱

    # 划分训练集和测试集
    testset_numer = int(len(images_filename) * test_frac) # 测试集图像个数
    testset_images = images_filename[:testset_numer]      # 获取拟移动至 test 目录的测试集图像文件名
    trainset_images = images_filename[testset_numer:]     # 获取拟移动至 train 目录的训练集图像文件名

    # 移动图像至 test 目录
    for image in testset_images:
        old_img_path = os.path.join(dataset_path, fruit, image)         # 获取原始文件路径
        new_test_path = os.path.join(dataset_path, 'val', fruit, image) # 获取 test 目录的新文件路径
        shutil.move(old_img_path, new_test_path) # 移动文件

    # 移动图像至 train 目录
    for image in trainset_images:
        old_img_path = os.path.join(dataset_path, fruit, image)           # 获取原始文件路径
        new_train_path = os.path.join(dataset_path, 'train', fruit, image) # 获取 train 目录的新文件路径
        shutil.move(old_img_path, new_train_path) # 移动文件
    
    # 删除旧文件夹
    assert len(os.listdir(old_dir)) == 0 # 确保旧文件夹中的所有图像都被移动走
    shutil.rmtree(old_dir) # 删除文件夹
    
    # 工整地输出每一类别的数据个数
    print('{:^18} {:^18} {:^18}'.format(fruit, len(trainset_images), len(testset_images)))
    
    # 保存到表格中
    df = df.append({'class':fruit, 'trainset':len(trainset_images), 'testset':len(testset_images)}, ignore_index=True)

# 重命名数据集文件夹
shutil.move(dataset_path, dataset_name+'_split')

# 数据集各类别数量统计表格，导出为 csv 文件
df['total'] = df['trainset'] + df['testset']
df.to_csv('数据量统计.csv', index=False)

df

小结

比较菜，写的不好，或哪里不对见谅，又不会的可以私聊我。只用改一一下文件路劲就行了。