| 11 | |
| 12 | |
| 13 | def read_split_data(root: str, val_rate: float = 0.2): |
| 14 | random.seed(0) # 保证随机结果可复现 |
| 15 | assert os.path.exists(root), "dataset root: {} does not exist.".format(root) |
| 16 | |
| 17 | # 遍历文件夹，一个文件夹对应一个类别 |
| 18 | flower_class = [cla for cla in os.listdir(root) if os.path.isdir(os.path.join(root, cla))] |
| 19 | # 排序，保证各平台顺序一致 |
| 20 | flower_class.sort() |
| 21 | # 生成类别名称以及对应的数字索引 |
| 22 | class_indices = dict((k, v) for v, k in enumerate(flower_class)) |
| 23 | json_str = json.dumps(dict((val, key) for key, val in class_indices.items()), indent=4) |
| 24 | with open('class_indices.json', 'w') as json_file: |
| 25 | json_file.write(json_str) |
| 26 | |
| 27 | train_images_path = [] # 存储训练集的所有图片路径 |
| 28 | train_images_label = [] # 存储训练集图片对应索引信息 |
| 29 | val_images_path = [] # 存储验证集的所有图片路径 |
| 30 | val_images_label = [] # 存储验证集图片对应索引信息 |
| 31 | every_class_num = [] # 存储每个类别的样本总数 |
| 32 | supported = [".jpg", ".JPG", ".png", ".PNG"] # 支持的文件后缀类型 |
| 33 | # 遍历每个文件夹下的文件 |
| 34 | for cla in flower_class: |
| 35 | cla_path = os.path.join(root, cla) |
| 36 | # 遍历获取supported支持的所有文件路径 |
| 37 | images = [os.path.join(root, cla, i) for i in os.listdir(cla_path) |
| 38 | if os.path.splitext(i)[-1] in supported] |
| 39 | # 排序，保证各平台顺序一致 |
| 40 | images.sort() |
| 41 | # 获取该类别对应的索引 |
| 42 | image_class = class_indices[cla] |
| 43 | # 记录该类别的样本数量 |
| 44 | every_class_num.append(len(images)) |
| 45 | # 按比例随机采样验证样本 |
| 46 | val_path = random.sample(images, k=int(len(images) * val_rate)) |
| 47 | |
| 48 | for img_path in images: |
| 49 | if img_path in val_path: # 如果该路径在采样的验证集样本中则存入验证集 |
| 50 | val_images_path.append(img_path) |
| 51 | val_images_label.append(image_class) |
| 52 | else: # 否则存入训练集 |
| 53 | train_images_path.append(img_path) |
| 54 | train_images_label.append(image_class) |
| 55 | |
| 56 | print("{} images were found in the dataset.".format(sum(every_class_num))) |
| 57 | print("{} images for training.".format(len(train_images_path))) |
| 58 | print("{} images for validation.".format(len(val_images_path))) |
| 59 | assert len(train_images_path) > 0, "number of training images must greater than 0." |
| 60 | assert len(val_images_path) > 0, "number of validation images must greater than 0." |
| 61 | |
| 62 | plot_image = False |
| 63 | if plot_image: |
| 64 | # 绘制每种类别个数柱状图 |
| 65 | plt.bar(range(len(flower_class)), every_class_num, align='center') |
| 66 | # 将横坐标0,1,2,3,4替换为相应的类别名称 |
| 67 | plt.xticks(range(len(flower_class)), flower_class) |
| 68 | # 在柱状图上添加数值标签 |
| 69 | for i, v in enumerate(every_class_num): |
| 70 | plt.text(x=i, y=v + 5, s=str(v), ha='center') |