def load_images_and_labels(dataset_dir, image_size): images = [] labels = [] class_labels = os.listdir(dataset_dir) for i, class_label in enumerate(class_labels): class_dir = os.path.join(dataset_dir, class_label) for image_file in os.listdir(class_dir): image_path = os.path.join(class_dir, image_file) image = Image.open(image_path).convert('RGB') image = image.resize(image_size) image = np.array(image) images.append(image) labels.append(i) images = np.array(images) labels = np.array(labels) return images, labels
时间: 2024-03-29 21:36:40 浏览: 73
这段代码的作用是什么?
这段代码定义了一个函数load_images_and_labels,用于加载指定目录下的图像和标签数据。函数的参数包括dataset_dir和image_size,其中dataset_dir表示图像数据所在的目录,image_size表示图像的尺寸。函数返回两个numpy数组,分别为images和labels,分别表示图像数据和对应的标签数据。
具体实现过程为,首先通过os.listdir函数获取dataset_dir下的所有子目录,每个子目录表示一个类别。然后遍历每个子目录,读取该目录下的所有图像文件,并将其转换为RGB格式、调整大小为image_size,并转换为numpy数组。同时,将该图像对应的标签记录在labels数组中,并用整数表示类别。最后,将所有图像和标签数据分别存储在numpy数组images和labels中,并返回这两个数组。
相关问题
import os.path import gzip import pickle import os import numpy as np import urllib url_base = 'http://yann.lecun.com/exdb/mnist/' key_file = { 'train_img':'train-images-idx3-ubyte.gz', 'train_label':'train-labels-idx1-ubyte.gz', 'test_img':'t10k-images-idx3-ubyte.gz', 'test_label':'t10k-labels-idx1-ubyte.gz' } dataset_dir = os.path.dirname(os.path.abspath("_file_")) save_file = dataset_dir + "/mnist.pkl" train_num=60000 test_num=10000 img_dim=(1,28,28) img_size=784 def _download(file_name): file_path = dataset_dir+"/"+file_name if os.path.exists(file_path): return print("Downloading"+file_name+" ... ") urllib.request.urlretrieve(url_base + file_name,file_path) print("Done") def download_mnist(): for v in key_file.values(): _download(v) def _load_label(file_name): file_path = dataset_dir+ "/" +file_name print("Converting" + file_name +"to Numpy Array ...") with gzip.open(file_path,'rb') as f: labels = np.frombuffer(f.read(),np.uint8,offset=8) print("Done") return labels def _load_img(file_name): file_path=dataset_dir+"/"+file_name print("Converting"+file_name+"to Numpy Array ...") with gzip.open(file_path,'rb') as f: data = np.frombuffer(f.read(),np.uint8,offset=16) data = data.reshape(-1,img_size) print("Done") return data def _convert_numpy(): dataset = {} dataset['train_img'] = _load_img(key_file['train_img']) dataset['train_label'] = _load_label(key_file['train_label']) dataset['test_img'] = _load_img(key_file['test_img']) dataset['test_label'] = _load_label(key_file['test_label']) return dataset def init_mnist(): download_mnist() dataset = _convert_numpy() print("Creating pickle file ...") with open(save_file,'wb') as f: pickle.dump(dataset,f,-1) print("Done") if __name__ =='__main__': init_mnist()
这段代码是用于下载MNIST数据集并将数据集转换成Numpy数组格式的函数。MNIST数据集是一个手写数字识别数据集,包含了60000张训练图片和10000张测试图片。在函数中,首先定义了数据集的下载地址和四个文件的名称,然后定义了四个函数用于下载和转换数据集。其中,_load_label和_load_img函数用于将标签和图片数据转换成Numpy数组格式,并且在转换过程中使用了gzip库解压缩数据。_convert_numpy函数用于将四个Numpy数组合成一个字典类型的数据集。最后,init_mnist函数用于下载数据集并将数据集转换成Numpy数组格式,并使用pickle库将数据集保存到本地文件中。如果MNIST数据集已经下载并保存到本地文件中,则直接加载本地文件中的数据集。
class MyDataSet(Dataset): """自定义数据集""" def __init__(self, images_path: list, images_class: list, transform=None): self.images_path = images_path self.images_class = images_class self.transform = transform def __len__(self): return len(self.images_path) def __getitem__(self, item): img = Image.open(self.images_path[item]) # RGB为彩色图片,L为灰度图片 if img.mode != 'RGB': raise ValueError("image: {} isn't RGB mode.".format(self.images_path[item])) label = self.images_class[item] if self.transform is not None: img = self.transform(img) print(label) return img, label @staticmethod def collate_fn(batch): # 官方实现的default_collate可以参考 # https://github.com/pytorch/pytorch/blob/67b7e751e6b5931a9f45274653f4f653a4e6cdf6/torch/utils/data/_utils/collate.py images, labels = tuple(zip(*batch)) images = torch.stack(images, dim=0) labels = torch.as_tensor(np.array(labels[:, 0])) return images, labels
这段代码是一个自定义的 PyTorch 数据集类 `MyDataSet`,它继承自 PyTorch 内置的 `Dataset` 类。在 `__init__` 方法中,该类接收两个列表 `images_path` 和 `images_class`,分别代表图片的路径和对应的类别。还可以传入一个可选的 `transform` 参数,用于对图片进行预处理。在 `__getitem__` 方法中,该类会根据索引 `item` 加载对应路径的图片,并返回预处理后的图片张量和对应的类别标签。
这个类还定义了一个静态方法 `collate_fn`,它接收一个批次的数据,对其进行处理并返回。在这个方法中,首先使用 Python 的 `zip` 函数和 `*` 运算符将批次中的图片张量和对应的标签分别打包成两个元组。然后使用 `torch.stack` 函数将图片张量拼接成一个大的张量,`dim=0` 表示在第 0 维进行拼接。最后,代码使用 NumPy 的 `array` 函数将标签列表转换为 NumPy 数组,然后使用 `[:, 0]` 语法取出了每个元组中的第一个标签,最后使用 `torch.as_tensor` 函数将其转换为 PyTorch 张量。
这个自定义数据集类的实现比较简单,但需要注意一些细节,比如 `zip` 函数的使用,`labels[:, 0]` 的含义等等。
阅读全文