自定义数据集的处理

宝可梦数据集

一共有五个文件，文件里面有不同的图片文件，不同格式，五个文件对应不同的宝可梦。

一、框架

class pokemon类下

1. def __ init __()
   三个参数：文件地址路径，大小resize，mode模型

    def __init__(self, root, resize, mode):
        super(Pokemon, self).__init__()

        self.root = root
        self.resize = resize

下面这段代码是为了将物种宝可梦的文件名转成0~5的label，首先创建一个空的字典。然后遍历这个文件里的所有文件名，这个字典的长度就是label。通过self.name2label[name] = len(self.name2label.keys())就可以写成 {‘pikachu’:3}的这种形式。下面就是取出images和labels。需要自己写一个load _csv函数。

self.name2label = {} # "sq...":0
        for name in sorted(os.listdir(os.path.join(root))):
            if not os.path.isdir(os.path.join(root, name)):
                continue

            self.name2label[name] = len(self.name2label.keys())

        # print(self.name2label)

        # image, label
        self.images, self.labels = self.load_csv('images.csv')

这一段代码作用是将数据集分割成60%的训练集，20%的验证集和20%的测试集。通过简单的切片操作就可以做到。

        if mode=='train': # 60%
            self.images = self.images[:int(0.6*len(self.images))]
            self.labels = self.labels[:int(0.6*len(self.labels))]
        elif mode=='val': # 20% = 60%->80%
            self.images = self.images[int(0.6*len(self.images)):int(0.8*len(self.images))]
            self.labels = self.labels[int(0.6*len(self.labels)):int(0.8*len(self.labels))]
        else: # 20% = 80%->100%
            self.images = self.images[int(0.8*len(self.images)):]
            self.labels = self.labels[int(0.8*len(self.labels)):]

2. def __ len __()

    def __len__(self):

        return len(self.images)

3. def __ getitem __()
   用来获取images和labels的idx，这边另外加上了一个标准化的过程。用torchvision里面的transforms来进行图片数据的一些操作。

    def __getitem__(self, idx):
        # idx~[0~len(images)]
        # self.images, self.labels
        # img: 'pokemon\\bulbasaur\\00000000.png'
        # label: 0
        img, label = self.images[idx], self.labels[idx]

        tf = transforms.Compose([
            lambda x:Image.open(x).convert('RGB'), # string path= > image data
            transforms.Resize((int(self.resize*1.25), int(self.resize*1.25))),
            transforms.RandomRotation(15),
            transforms.CenterCrop(self.resize),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ])

        img = tf(img)
        label = torch.tensor(label)


        return img, label

4. def load_csv()
   传入的参数为csv文件名
   首先就是外面包裹着的if not语句，用来判断文件夹中是否存在这样一个csv文件，如果存在，那么将跳过创建的这一部分，直接执行下面的打开文件夹将csv文件中的img和label都写入images和label的列表的操作。
   然后创建空列表images，遍历name2label字典中的键，将所有图片文件传入列表当中，glob.glob的作用就是遍历所有路径中的文件。然后将images的img和label都写入到csv文件中去。得到label的话首先要得到name，就通过name = img.split(os.sep)[-2]来得到。
   第二部分就是从csv文件中读出img和label信息到新的列表当中。加上个断言语句确保img和labels的数量一样。

    def load_csv(self, filename):

        if not os.path.exists(os.path.join(self.root, filename)):
            images = []
            for name in self.name2label.keys():
                # 'pokemon\\mewtwo\\00001.png
                images += glob.glob(os.path.join(self.root, name, '*.png'))
                images += glob.glob(os.path.join(self.root, name, '*.jpg'))
                images += glob.glob(os.path.join(self.root, name, '*.jpeg'))

            # 1167, 'pokemon\\bulbasaur\\00000000.png'
            print(len(images), images)

            random.shuffle(images)
            with open(os.path.join(self.root, filename), mode='w', newline='') as f:
                writer = csv.writer(f)
                for img in images: # 'pokemon\\bulbasaur\\00000000.png'
                    name = img.split(os.sep)[-2]
                    label = self.name2label[name]
                    # 'pokemon\\bulbasaur\\00000000.png', 0
                    writer.writerow([img, label])
                print('writen into csv file:', filename)

        # read from csv file
        images, labels = [], []
        with open(os.path.join(self.root, filename)) as f:
            reader = csv.reader(f)
            for row in reader:
                # 'pokemon\\bulbasaur\\00000000.png', 0
                img, label = row
                label = int(label)

                images.append(img)
                labels.append(label)

        assert len(images) == len(labels)

        return images, labels

5. def denormalize()
   这个代码是用来denormailze的，因为之前normalize过了，把数据弄到-1~1之间了，所以用vizdom可视化的时候会看不清楚。其实就是按照normalize的公式反过来用来一边，但是要注意维度要一致。

    def denormalize(self, x_hat):

        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]

        # x_hat = (x-mean)/std
        # x = x_hat*std + mean
        # x: [c, h, w]
        # mean: [3] => [3, 1, 1]
        mean = torch.tensor(mean).unsqueeze(1).unsqueeze(1)
        std = torch.tensor(std).unsqueeze(1).unsqueeze(1)
        # print(mean.shape, std.shape)
        x = x_hat * std + mean

        return x

6. def _ main _()
   实例化db, iter()是个迭代器，可以遍历db中的元素，next就用来一个一个输出这些元素。viz.image()可以可视化过程。Dataloader用来把所有数据加载进去，实例化为loader。

def main():

    import  visdom
    import  time
    import  torchvision

    viz = visdom.Visdom()

    # tf = transforms.Compose([
    #                 transforms.Resize((64,64)),
    #                 transforms.ToTensor(),
    # ])
    # db = torchvision.datasets.ImageFolder(root='pokemon', transform=tf)
    # loader = DataLoader(db, batch_size=32, shuffle=True)
    #
    # print(db.class_to_idx)
    #
    # for x,y in loader:
    #     viz.images(x, nrow=8, win='batch', opts=dict(title='batch'))
    #     viz.text(str(y.numpy()), win='label', opts=dict(title='batch-y'))
    #
    #     time.sleep(10)


    db = Pokemon('pokemon', 64, 'train')

    x,y = next(iter(db))
    print('sample:', x.shape, y.shape, y)

    viz.image(db.denormalize(x), win='sample_x', opts=dict(title='sample_x'))

    loader = DataLoader(db, batch_size=32, shuffle=True, num_workers=8)

    for x,y in loader:
        viz.images(db.denormalize(x), nrow=8, win='batch', opts=dict(title='batch'))
        viz.text(str(y.numpy()), win='label', opts=dict(title='batch-y'))

        time.sleep(10)