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)
