代码一: 按照比例划分数据集
import os
import random
import json
def read_split_data(root: str,test_rate: float = 0.2):
"""
对图片数据集进行分割
:param root: 数据集所在的路径(不同类型图片所在文件夹路径)
:param test_rate: 测试集在数据集中所占的比例
:return: 训练图片路径,训练图片标签,测试集图片路径,测试集图片标签
"""
random.seed(0) # 保证随机结果可复现
assert os.path.exists(root), "dataset root: {} does not exist.".format(root)
# 遍历文件夹,一个文件夹对应一个类别
texture_class = [cla for cla in os.listdir(root)
if os.path.isdir(os.path.join(root, cla))]
# 排序,保证顺序一致
texture_class.sort()
# 生成类别名称以及对应的数字索引
class_indices = dict((k, v) for v, k in enumerate(texture_class))
json_str = json.dumps(dict((val, key) for key, val in class_indices.items()), indent=4)
with open(root+'/class_indices.json', 'w') as json_file:
json_file.write(json_str)
train_images_path = [] # 存储训练集的所有图片路径
train_images_label = [] # 存储训练集图片对应索引信息
test_images_path = [] # 存储测试集的所有图片路径
test_images_label = [] # 存储测试集图片对应索引信息
every_class_num = [] # 存储每个类别的样本总数
supported = [".jpg", ".JPG", ".png", ".PNG"] # 支持的文件后缀类型
# 遍历每个文件夹下的文件
for cla in texture_class:
cla_path = os.path.join(root, cla)
# 遍历获取supported支持的所有文件路径
images = [os.path.join(root, cla, i) for i in os.listdir(cla_path)
if os.path.splitext(i)[-1] in supported]
# 获取该类别对应的索引
image_class = class_indices[cla]
# 记录该类别的样本数量
every_class_num.append(len(images))
# 按比例随机采样验证样本
random.shuffle(images)
test_path = images[: int(len(images) * test_rate)]
# test_path = images[int(len(images) * val_rate):int(len(images) * (test_rate+val_rate))]
for img_path in images:
if img_path in test_path: # 如果该路径在采样的验证集样本中则存入验证集
test_images_path.append(img_path)
test_images_label.append(image_class)
else: # 否则存入训练集
train_images_path.append(img_path)
train_images_label.append(image_class)
print("{} images were found in the .".format(sum(every_class_num)))
print("{} images for training.".format(len(train_images_path)))
print("{} images for testing.".format(len(test_images_path)))
代码二:画图显示每个类别数量的柱状图
import matplotlib.pyplot as plt
plot_image = True
if plot_image:
# 绘制每种类别个数柱状图
plt.bar(range(len(_class)), every_class_num, align='center')
# 将横坐标0,1,2,3,4替换为相应的类别名称
plt.xticks(range(len(texture_class)), texture_class)
# 在柱状图上添加数值标签
for i, v in enumerate(every_class_num):
plt.text(x=i, y=v + 5, s=str(v), ha='center')
# 设置x坐标
plt.xlabel('image class')
# 设置y坐标
plt.ylabel('number of images')
# 设置柱状图的标题
plt.title('texture class distribution')
plt.show()
代码三:全部代码
import os
import random
import json
def read_split_data(root: str,test_rate: float = 0.2):
"""
对图片数据集进行分割
:param root: 数据集所在的路径(不同类型图片所在文件夹路径)
:param test_rate: 测试集在数据集中所占的比例
:return: 训练图片路径,训练图片标签,测试集图片路径,测试集图片标签
"""
random.seed(0) # 保证随机结果可复现
assert os.path.exists(root), "dataset root: {} does not exist.".format(root)
# 遍历文件夹,一个文件夹对应一个类别
texture_class = [cla for cla in os.listdir(root)
if os.path.isdir(os.path.join(root, cla))]
# 排序,保证顺序一致
texture_class.sort()
# 生成类别名称以及对应的数字索引
class_indices = dict((k, v) for v, k in enumerate(texture_class))
json_str = json.dumps(dict((val, key) for key, val in class_indices.items()), indent=4)
with open(root+'/class_indices.json', 'w') as json_file:
json_file.write(json_str)
train_images_path = [] # 存储训练集的所有图片路径
train_images_label = [] # 存储训练集图片对应索引信息
test_images_path = [] # 存储测试集的所有图片路径
test_images_label = [] # 存储测试集图片对应索引信息
every_class_num = [] # 存储每个类别的样本总数
supported = [".jpg", ".JPG", ".png", ".PNG"] # 支持的文件后缀类型
# 遍历每个文件夹下的文件
for cla in texture_class:
cla_path = os.path.join(root, cla)
# 遍历获取supported支持的所有文件路径
images = [os.path.join(root, cla, i) for i in os.listdir(cla_path)
if os.path.splitext(i)[-1] in supported]
# 获取该类别对应的索引
image_class = class_indices[cla]
# 记录该类别的样本数量
every_class_num.append(len(images))
# 按比例随机采样验证样本
random.shuffle(images)
test_path = images[: int(len(images) * test_rate)]
# test_path = images[int(len(images) * val_rate):int(len(images) * (test_rate+val_rate))]
for img_path in images:
if img_path in test_path: # 如果该路径在采样的验证集样本中则存入验证集
test_images_path.append(img_path)
test_images_label.append(image_class)
else: # 否则存入训练集
train_images_path.append(img_path)
train_images_label.append(image_class)
plot_image = True
if plot_image:
# 绘制每种类别个数柱状图
plt.bar(range(len(_class)), every_class_num, align='center')
# 将横坐标0,1,2,3,4替换为相应的类别名称
plt.xticks(range(len(texture_class)), texture_class)
# 在柱状图上添加数值标签
for i, v in enumerate(every_class_num):
plt.text(x=i, y=v + 5, s=str(v), ha='center')
# 设置x坐标
plt.xlabel('image class')
# 设置y坐标
plt.ylabel('number of images')
# 设置柱状图的标题
plt.title('texture class distribution')
plt.show()
print("{} images were found in the .".format(sum(every_class_num)))
print("{} images for training.".format(len(train_images_path)))
print("{} images for testing.".format(len(test_images_path)))