LoadDataST.py

#!/usr/bin/env python
# coding: utf-8

# In[ ]:


import os
import torch
from torch.utils.data import Dataset
from torchvision.transforms import transforms
import numpy as np
import collections
from PIL import Image
import csv
import random


class ImagenetMini(Dataset):
    """
    put mini-imagenet files as :
    root :
        |- images/*.jpg includes all imgeas
        |- train.csv
        |- test.csv
        |- val.csv
    NOTICE: meta-learning is different from general supervised learning, especially the concept of batch and set.
    batch: contains several sets
    sets: conains n_way * k_shot for meta-train set, n_way * n_query for meta-test set.
    """

    def __init__(self, root, mode, batchsz, resize, startidx=0):
        """
        :param root: root path of mini-imagenet
        :param mode: train, val or test
        :param batchsz: batch size of sets, not batch of imgs
        :param n_way:
        :param k_shot:
        :param k_query: num of qeruy imgs per class
        :param resize: resize to
        :param startidx: start to index label from startidx
        """

        self.batchsz = batchsz  #batch of imgs
        self.resize = resize  # resize to
        self.startidx = startidx  # index label not from 0, but from startidx
        print('shuffle DB :%s, b:%d, resize:%d' % (
        mode, batchsz, resize))

        if mode == 'train':
#             self.transform = transforms.Compose([lambda x: Image.open(x).convert('RGB'),
#                                                  transforms.Resize((self.resize, self.resize)),
#                                                  # transforms.RandomHorizontalFlip(),
#                                                  # transforms.RandomRotation(5),
#                                                  transforms.ToTensor(),
#                                                  transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
#                                                  ])
            self.transform = transforms.Compose([lambda x: Image.open(x).convert('RGB'),
                                                 transforms.Resize((self.resize, self.resize)),
                                                 # transforms.RandomHorizontalFlip(),
                                                 # transforms.RandomRotation(5),
                                                 transforms.ToTensor(),
                                                 transforms.Normalize((0, 0, 0), (1, 1, 1))
                                                 ])
        else:
#             self.transform = transforms.Compose([lambda x: Image.open(x).convert('RGB'),
#                                                  transforms.Resize((self.resize, self.resize)),
#                                                  transforms.ToTensor(),
#                                                  transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
#                                                  ])
            self.transform = transforms.Compose([lambda x: Image.open(x).convert('RGB'),
                                                 transforms.Resize((self.resize, self.resize)),
                                                 transforms.ToTensor(),
                                                 transforms.Normalize((0, 0, 0), (1, 1, 1))
                                                 ])

        self.path = os.path.join(root, 'images')  # image path
        csvdata = self.loadCSV(os.path.join(root, mode + '.csv'))  # csv path
        self.data = []
        self.img2label = {}
        for i, (k, v) in enumerate(csvdata.items()):
            self.data.extend(v)  # [[img1, img2, ...], [img111, ...]]
            self.img2label[k] = i + self.startidx  # {"img_name[:9]":label}
        self.cls_num = len(self.data)

#         self.create_batch(self.batchsz)
        
        
        support_x = torch.FloatTensor(self.cls_num, 3, self.resize, self.resize)
        support_x_temp = np.array(self.data)#.tolist()
        #support_y = np.zeros((self.cls_num), dtype=np.int)
        flatten_x = [os.path.join(self.path, item) for item in support_x_temp]
        temp = [self.img2label[item[:9]] for item in support_x_temp]
        support_y = np.array(temp).astype(np.int32)
                             
        for i, path in enumerate(flatten_x):
            support_x[i] = self.transform(path)

        self.loading_data = DataSubset(support_x, torch.LongTensor(support_y))

    def loadCSV(self, csvf):
        """
        return a dict saving the information of csv
        :param splitFile: csv file name
        :return: {label:[file1, file2 ...]}
        """
        dictLabels = {}
        with open(csvf) as csvfile:
            csvreader = csv.reader(csvfile, delimiter=',')
            next(csvreader, None)  # skip (filename, label)
            for i, row in enumerate(csvreader):
                filename = row[0]
                label = row[1]
                # append filename to current label
                if label in dictLabels.keys():
                    dictLabels[label].append(filename)
                else:
                    dictLabels[label] = [filename]
        return dictLabels

#     def create_batch(self, batchsz):
#         """
#         create batch for meta-learning.
#         ×episode× here means batch, and it means how many sets we want to retain.
#         :param episodes: batch size
#         :return:
#         """
#         self.x_batch = []  # support set batch
#         for b in range(batchsz):  # for each batch
#             # 1.select n_way classes randomly
#             selected_cls = np.random.choice(self.cls_num, self.n_way, False)  # no duplicate
#             np.random.shuffle(selected_cls)
#             support_x = []
#             query_x = []
#             for cls in selected_cls:
#                 # 2. select k_shot + k_query for each class
#                 selected_imgs_idx = np.random.choice(len(self.data[cls]), self.k_shot + self.k_query, False)
#                 np.random.shuffle(selected_imgs_idx)
#                 indexDtrain = np.array(selected_imgs_idx[:self.k_shot])  # idx for Dtrain
#                 indexDtest = np.array(selected_imgs_idx[self.k_shot:])  # idx for Dtest
#                 support_x.append(
#                     np.array(self.data[cls])[indexDtrain].tolist())  # get all images filename for current Dtrain
#                 query_x.append(np.array(self.data[cls])[indexDtest].tolist())

#             # shuffle the correponding relation between support set and query set
#             random.shuffle(support_x)
#             random.shuffle(query_x)

#             self.support_x_batch.append(support_x)  # append set to current sets
#             self.query_x_batch.append(query_x)  # append sets to current sets

#     def __getitem__(self, index):
#         """
#         index means index of sets, 0<= index <= batchsz-1
#         :param index:
#         :return:
#         """
#         # [setsz, 3, resize, resize]
#         support_x = torch.FloatTensor(self.setsz, 3, self.resize, self.resize)
#         # [setsz]
#         support_y = np.zeros((self.setsz), dtype=np.int)
#         # [querysz, 3, resize, resize]
#         query_x = torch.FloatTensor(self.querysz, 3, self.resize, self.resize)
#         # [querysz]
#         query_y = np.zeros((self.querysz), dtype=np.int)

#         flatten_support_x = [os.path.join(self.path, item)
#                              for sublist in self.support_x_batch[index] for item in sublist]
#         support_y = np.array(
#             [self.img2label[item[:9]]  # filename:n0153282900000005.jpg, the first 9 characters treated as label
#              for sublist in self.support_x_batch[index] for item in sublist]).astype(np.int32)

#         flatten_query_x = [os.path.join(self.path, item)
#                            for sublist in self.query_x_batch[index] for item in sublist]
#         query_y = np.array([self.img2label[item[:9]]
#                             for sublist in self.query_x_batch[index] for item in sublist]).astype(np.int32)

#         # print('global:', support_y, query_y)
#         # support_y: [setsz]
#         # query_y: [querysz]
#         # unique: [n-way], sorted
#         unique = np.unique(support_y)
#         random.shuffle(unique)
#         # relative means the label ranges from 0 to n-way
#         support_y_relative = np.zeros(self.setsz)
#         query_y_relative = np.zeros(self.querysz)
#         for idx, l in enumerate(unique):
#             support_y_relative[support_y == l] = idx
#             query_y_relative[query_y == l] = idx

#         # print('relative:', support_y_relative, query_y_relative)

#         for i, path in enumerate(flatten_support_x):
#             support_x[i] = self.transform(path)

#         for i, path in enumerate(flatten_query_x):
#             query_x[i] = self.transform(path)
#         # print(support_set_y)
#         # return support_x, torch.LongTensor(support_y), query_x, torch.LongTensor(query_y)

#         return support_x, torch.LongTensor(support_y_relative), query_x, torch.LongTensor(query_y_relative)

#     def __len__(self):
#         # as we have built up to batchsz of sets, you can sample some small batch size of sets.
#         return self.batchsz

    
    
class DataSubset(object):
    def __init__(self, xs, ys, num_examples=None, seed=None):

        if seed is not None:
            np.random.seed(99)
        self.xs = xs
        self.n = len(xs)
        self.ys = ys
        self.batch_start = 0
        self.cur_order = np.random.permutation(self.n)

    def get_next_batch(self, batch_size, multiple_passes=False, reshuffle_after_pass=True):
        # np.random.seed(99)
        if self.n < batch_size:
            raise ValueError('Batch size can be at most the dataset size')
        if not multiple_passes:
            actual_batch_size = min(batch_size, self.n - self.batch_start)
            if actual_batch_size <= 0:
                raise ValueError('Pass through the dataset is complete.')
            batch_end = self.batch_start + actual_batch_size
            batch_xs = self.xs[self.cur_order[self.batch_start : batch_end], ...]
            batch_ys = self.ys[self.cur_order[self.batch_start : batch_end], ...]
            self.batch_start += actual_batch_size
            return batch_xs, batch_ys
        actual_batch_size = min(batch_size, self.n - self.batch_start)
        if actual_batch_size < batch_size:
            if reshuffle_after_pass:
                self.cur_order = np.random.permutation(self.n)
            self.batch_start = 0
        batch_end = self.batch_start + batch_size
#         batch_xs = self.xs[self.cur_order[self.batch_start : batch_end], ...]
#         batch_ys = self.ys[self.cur_order[self.batch_start : batch_end], ...]

        batch_xs = self.xs[self.cur_order[self.batch_start : batch_end], ...]
        batch_ys = self.ys[self.cur_order[self.batch_start : batch_end], ...]
       

        self.batch_start += actual_batch_size
        return batch_xs, batch_ys


if __name__ == '__main__':
    # the following episode is to view one set of images via tensorboard.
    from torchvision.utils import make_grid
    from matplotlib import pyplot as plt
    from tensorboardX import SummaryWriter
    import time

    plt.ion()

    tb = SummaryWriter('runs', 'mini-imagenet')
    mini = MiniImagenet('../../../dataset/', mode='train', batchsz=1000, resize=168)

    for i, set_ in enumerate(mini):
        # support_x: [k_shot*n_way, 3, 84, 84]
        support_x, support_y, query_x, query_y = set_

        support_x = make_grid(support_x, nrow=2)
        query_x = make_grid(query_x, nrow=2)

        plt.figure(1)
        plt.imshow(support_x.transpose(2, 0).numpy())
        plt.pause(0.5)
        plt.figure(2)
        plt.imshow(query_x.transpose(2, 0).numpy())
        plt.pause(0.5)

        tb.add_image('support_x', support_x)
        tb.add_image('query_x', query_x)

        time.sleep(5)

    tb.close()