pytorch實現花朵數據集讀取

import os
from PIL import Image
from torch.utils import data
import numpy as np
from torchvision import transforms as T


class My_Data(data.Dataset):

    def __init__(self, root, transforms=None, train=True, test=False):
        '''
        目標：獲取所有圖片路徑，並根據訓練、驗證、測試划分數據
        '''
        self.test = test
        classs = os.listdir(root)
        imgs = []
        labels = []
        for idx, folder in enumerate(classs):
            cate = os.path.join(root, folder)
            for img_num, im in enumerate(os.listdir(cate)):
                img_path = os.path.join(cate, im)
                #打包圖片路徑（轉換為list）
                imgs.append(img_path)
                #打包標簽路徑（轉換為list）
                labels.append(idx)
        if self.test:
            imgs = sorted(imgs, key=lambda x: int(x.split('.')[-2].split('/')[-1]))
        else:

            imgs = list(zip(imgs , labels))
            #將圖片路徑與標簽打包成一個list

        imgs_num = len(imgs)

        # shuffle imgs
        np.random.seed(100)
        imgs = np.random.permutation(imgs)

        # 划分訓練、驗證集，驗證:訓練 = 3:7
        if self.test:
            self.imgs = imgs
        elif train:
            self.imgs = imgs[:int(0.7 * imgs_num)]
        else:
            self.imgs = imgs[int(0.7 * imgs_num):]

        if transforms is None:

            # 數據轉換操作，測試驗證和訓練的數據轉換有所區別
            normalize = T.Normalize(mean=[0.485, 0.456, 0.406],
                                    std=[0.229, 0.224, 0.225])

            # 測試集和驗證集不用數據增強
            if self.test or not train:
                self.transforms = T.Compose([
                    T.Resize(32),
                    T.CenterCrop(32),
                    T.ToTensor(),
                    normalize
                ])
                # 訓練集需要數據增強
            else:
                self.transforms = T.Compose([
                    T.Resize(32),
                    T.RandomResizedCrop(32),
                    T.RandomHorizontalFlip(),
                    T.ToTensor(),
                    normalize
                ])

    def __getitem__(self,index):
        '''
        返回一張圖片的數據
        對於測試集，沒有label，返回圖片id，如1000.jpg返回1000
        送入一個batch_size的數據
        '''

        img_lables = self.imgs[index]
        img_path = img_lables[0]

        if self.test:
            label = int(self.imgs[index].split('.')[-2].split('/')[-1])
        else:
            label = int(img_lables[1])

        data = Image.open(img_path)
        data = self.transforms(data)
        return data, label

    def __len__(self):
        '''
        返回數據集中所有圖片的個數
        '''
        return len(self.imgs)

作為備份使用。