生成对抗网络——dataset对文件进行读入

近期在进行生成对抗网络的实验课，对于生成对抗网络(GAN)，最重要的就是确定训练集和各项参数。其中对于训练集的处理，即对图片的处理，有很多种方式，我这里写了一个myDataset类对图片进行输入处理

对于训练集，写了一个读入函数myDataset，对选定的训练集进行读入

import numpy as np
import torch
import os
import random
from PIL import Image
from torch.utils.data import Dataset

class myDataset(Dataset):
    def __init__(self, data_dir, transform):

        self.data_dir = data_dir
        self.transform = transform
        self.img_names = [name for name in list(filter(lambda x: x.endswith(".jpg"), os.listdir(self.data_dir)))]

    def __getitem__(self, index):
        path_img = os.path.join(self.data_dir, self.img_names[index])
        img = Image.open(path_img).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)

        return img

    def __len__(self):
        if len(self.img_names) == 0:
            raise Exception("\ndata_dir:{} is a empty dir! Please checkout your path to images!".format(self.data_dir))
        return len(self.img_names)

把文件命名为my_dataset.py文件，保存在:Python的Lib\site-packages\tools目录下