本文为看了小土堆的pytorch的教学视频后的个人总结。视频链接。
本文用到的数据集下载地址: 数据下载地址 密码:5suq
- 首先我们先导入一些之后操作需要的模块。
from torch.utils.data import Dataset
from PIL import Image
import os
数据集合和代码的相对位置:
class MyData(Dataset):
def __init__(self, root_dir, label_dir):
self.root_dir = root_dir
self.label_dir = label_dir
self.path = os.path.join(self.root_dir, self.label_dir)
self.img_path = os.listdir(self.path)
def __getitem__(self, idx):
img_name = self.img_path[idx]
img_item_path = os.path.join(self.root_dir, self.label_dir, img_name)
img = Image.open(img_item_path)
label = self.label_dir
return img, label
def __len__(self):
return len(self.img_path)
os.path.join
对路径进行拼接。
os.listdir()
返回指定的文件夹包含的文件名字的列表。
root_dir = "./data/train"
label_dir = "ants"
ants_dataset = MyData(root_dir, label_dir)
img, label = ants_dataset[1]
print(label)
img.show()
这样操作之后便可以打开相应文件夹中的图片。