PyTorch1加载数据初认识+Dataset

 1 from torch.utils.data import Dataset
 2 from PIL import Image
 3 
 4 import cv2
 5 import os
 6 
 7 class MyData(Dataset):
 8     def __init__(self, root_dir, lable_dir):
 9         self.root_dir = root_dir
10         self.lable_dir = lable_dir
11         self.path = os.path.join(self.root_dir, self.lable_dir)
12         self.img_path = os.listdir(self.path)
13 
14     def __getitem__(self, idx):
15         img_name = self.img_path[idx]
16         img_item_path = os.path.join(self.root_dir, self.lable_dir, img_name)
17         img = Image.open(img_item_path)
18         lable = self.lable_dir
19         return img, lable
20 
21     def __len__(self):
22         return len(self.img_path)
23 
24 root_dir = "E:\\python\\learn_pytorch\\dataset\\train"
25 ants_lable_dir = "ants"
26 bees_lable_dir = "bees"
27 ants_dataset = MyData(root_dir, ants_lable_dir)
28 bees_dataset = MyData(root_dir, bees_lable_dir)

标签：lable,img,PyTorch1,self,Dataset,path,root,dir,加载
来源： https://www.cnblogs.com/0xiaoyu/p/16119206.html