上一篇文章,我们用了莫凡大佬的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)