一、框架

(一)快速获得API(#以Dataset为例)

1
2
3
help(Dataset)
Dataset??
或者直接Ctrl+Dataset

获取图像的库Pillow

Python 学习笔记之—— PIL 库 - 知乎 (zhihu.com)

文件操作OS

简单明了!Python中的OS模块 - 知乎 (zhihu.com)

(二)Datase

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
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,index):
img_name = self.img_path[index]
img_item_name = os.path.join(self.root_dir, self.label_dir,img_name)
img = Image.open(img_item_name)
label = self.label_dir
return img,label

def __len__(self):
return len(self.img_path)

root_dir = '/home/ycy523/test/dataset/train'
label='ants'
ants_dataset = MyData(root_dir, label)