生成一个根据老师提供的手机相册分类数据集,制作成训练图像数据集、测试图像数据集、训练标签数据集和测试标签数据集。将四部分分别保存成numpy文件的代码
时间: 2024-05-13 21:16:30 浏览: 112
以下是一个可能的代码实现,需要注意的是,具体的数据集和文件路径需要根据实际情况修改。该代码实现的功能为:将指定路径下的所有图片随机分为训练集和测试集,同时生成对应的标签,最后将四个文件保存为numpy文件。
```python
import os
import random
import numpy as np
from PIL import Image
# 图像文件路径
image_dir = './dataset'
# 训练集和测试集的比例
train_ratio = 0.8
# 随机种子
random_seed = 42
# 读取所有图像文件名
image_files = []
for root, dirs, files in os.walk(image_dir):
for file in files:
if file.endswith('.jpg') or file.endswith('.jpeg') or file.endswith('.png'):
image_files.append(os.path.join(root, file))
# 打乱顺序
random.seed(random_seed)
random.shuffle(image_files)
# 分离训练集和测试集
split_index = int(train_ratio * len(image_files))
train_files = image_files[:split_index]
test_files = image_files[split_index:]
# 生成标签
categories = set(os.listdir(image_dir))
category2id = {category: i for i, category in enumerate(categories)}
train_labels = [category2id[os.path.basename(os.path.dirname(file))] for file in train_files]
test_labels = [category2id[os.path.basename(os.path.dirname(file))] for file in test_files]
# 读取图像并转换为numpy数组
def read_images(image_files):
images = []
for file in image_files:
image = Image.open(file).convert('RGB')
image = np.array(image)
images.append(image)
return np.array(images)
# 保存为numpy文件
np.save('train_images.npy', read_images(train_files))
np.save('test_images.npy', read_images(test_files))
np.save('train_labels.npy', np.array(train_labels))
np.save('test_labels.npy', np.array(test_labels))
```
以上代码中,`image_dir` 变量定义了图像文件所在的路径,`train_ratio` 变量定义了训练集和测试集的比例,`random_seed` 定义了随机种子,用于打乱图像文件的顺序。随机打乱后,根据 `train_ratio` 分离出训练集和测试集,然后根据文件名生成标签,最后将图像和标签分别保存为numpy文件。
阅读全文