【pytorch】学会pytorch dataloader数据加载（二）---MINC实战

上一篇文章，我们用了莫凡大佬的pytorch dataloader的玩具例子，简单的讲解了一下pytorch的数据加载。总结来说，pytorch的数据加载主要就是涉及两方面，一个是构造Datas，一个是构造Dataloader。本文将根据MINC这个数据集进行实战。

MINC项目主页： http://opensurfaces.cs.cornell.edu/publications/minc/。
论文名称： Material Recognition in the Wild with the Materials in Context Database
MINC是对精细纹理的23类类别进行分类的一篇文章，发表在CVPR2015。

MINC数据集介绍：

本文主要介绍MINC-2005，这是一个图像块分类数据，一共23类，每类含2500样本，这是MINC的子数据集，每类样本是均匀取样，图像快大小在362x362。

MINC-2500 is a patch classification dataset with 2500 samples per category (Section 5.4 of the paper). This is a subset of MINC where samples have been sized to 362 x 362 and each category is sampled evenly. The original resolution images are not needed as we include the extracted patches in the archive.

下载下来其数据集：

其中images文件夹里：23类数据

labels文件夹里：各个split好的txt

先看Readme：就能基本明白数据是怎么回事。

categories.txt - An ordered list of categories (zero-indexed labels).
images/CATEGORY/*.jpg - patches for the corresponding category
labels/trainN.txt - List of filenames, N = 1, ..., 5 (different splits).
labels/validateN.txt - Same
labels/testN.txt - Same

MINC-2500的dataset：

一、 MINC-2500的dataset构造：

class MINC2500(Dataset):
    def __init__(self, root_dir, set_type, split, transform=None):
        self.root_dir = root_dir
        self.set_type = set_type
        self.transform = transform
        # Those values are computed using the script 'get_minc2500_norm.py'
        self.mean = Tensor([0.507207, 0.458292, 0.404162])
        self.std = Tensor([0.254254, 0.252448, 0.266003])

        # Get the material categories from the categories.txt file
        file_name = os.path.join(root_dir, 'categories.txt')
        self.categories = {}
        with open(file_name, 'r') as f:
            for i, line in enumerate(f):
                # The last line char (\n) must be removed
                self.categories[line[:-1]] = i

        # Load the image paths
        self.data = []
        file_name = os.path.join(root_dir, 'labels')
        # For the moment I use only the first split
        file_name = os.path.join(file_name, set_type + str(split) + '.txt')
        with open(file_name, 'r') as f:
            for line in f:
                img_path = line.split(os.sep)
                # The last line char (\n) must be removed
                self.data.append([line[:-1], self.categories[img_path[1]]])

    def __len__(self):
        '''返回数据长度'''
        return len(self.data)

    def __getitem__(self, idx):
        '''返回单次迭代的数据，经过transform进行处理后输出'''
        img_path = os.path.join(self.root_dir, self.data[idx][0])
        image = Image.open(img_path)
        # Sometimes the images are opened as grayscale, so I need to force RGB
        image = image.convert('RGB')
        if self.transform:
            image = self.transform(image)
            if is_tensor(image):
                image = transforms.Normalize(self.mean, self.std)(image)

        return image, self.data[idx][1]

二、Dataloader构造

    '''声明需要对数据进行操作的transforms， 传入Dataset中'''
    train_trans = transforms.Compose([
        transforms.RandomSizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor()
    ])
    val_trans = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor()
    ])

    train_set = MINC2500(root_dir=args.train_folder, set_type='train', split=1, transform=train_trans)
    val_set   = MINC2500(root_dir=args.train_folder, set_type='validate', split=1, transform=val_trans)
    
    '''设置dataloader，每个batch输出'''
    train_loader = DataLoader(dataset=train_set, batch_size=args.batchsize, shuffle=True,  num_workers=args.workers, pin_memory=True)
    val_loader   = DataLoader(dataset=val_set,   batch_size=args.batchsize, shuffle=False, num_workers=args.workers, pin_memory=True)