语义分割学习系列（七）FCN模型加载VOC格式数据集的代码分析

前言

无论是哪个语义分割算法模型，比如FCN或segnet等，都会相应的代码来读取、加载以及预处理不同格式的数据集。毕竟每个格式数据集的目录组织形式都不尽相同。

加载VOC数据集的代码分析

在train.py中，构建VOC train和val数据集对象的代码分别是：

train_data = voc_loader.VOC2012ClassSeg(root=data_path, split='train', transform=True)

val_data = voc_loader.VOC2012ClassSeg(root=data_path,
                            split='val',
                            transform=True)

下面以train_data为例来分析训练样本数据集的读取和预处理过程。

上面两个代码都调用的是voc_loader.py里面的class： VOC2012ClassSeg。所以我们到voc_loader.py来看该class的定义：

class VOC2012ClassSeg(VOCClassSegBase):
    def __init__(self, root, split='train', transform=False):
        super(VOC2012ClassSeg, self).__init__(
            root, split=split, transform=transform)

该类的构造函数调用super()，它会先调用其父类 VOCClassSegBase(root, split=split, transform=transform)

所以我们又来看VOCClassSegBase类的构造函数定义。

    def __init__(self, root, split='train', transform=True):
        self.root = root
        self.split = split
        self._transform = transform

        # VOC2011 and others are subset of VOC2012
        dataset_dir = osp.join(self.root, 'VOC/VOCdevkit/VOC2012')
        # dataset_dir = osp.join(self.root, 'VOC2007')

        self.files = collections.defaultdict(list)
        for split_file in ['train', 'val']:
            imgsets_file = osp.join(
                dataset_dir, 'ImageSets/Segmentation/%s.txt' % split_file)
            for img_name in open(imgsets_file):
                img_name = img_name.strip()
                img_file = osp.join(dataset_dir, 'JPEGImages/%s.jpg' % img_name)
                lbl_file = osp.join(dataset_dir, 'SegmentationClass/%s.png' % img_name)
                self.files[split_file].append({
                    'img': img_file,
                    'lbl': lbl_file,
                })

该构造函数会先对train_loader对象的成员变量：root，split以及_transform进行赋值，然后将files成员变量按'train'和'val‘来分别赋值img和lbl，这里img是JPEGImages里面的样本图片像素值，而lbl则来自SegmentationClass的png像素值。

这里有个注意点就是，虽然train_data和val_data的files成员都既包括train样本集又包括val样本集数据，但是它们的成员属性split则分别是'train'和'val'。 其实到这里样本数据集都已经读取到train_datah和val_data对象里面了。

接下来在train.py的train和test函数中 会分别enumerate 训练数据和验证数据集。我们这里只看train数据集的加载，它是按batch size里逐批预处理和加载的。在train.py中，相关函数为：

train_loader = torch.utils.data.DataLoader(train_data,
                                           batch_size=batch_size,
                                           shuffle=True,
                                           num_workers=5)

... ...

def train(epoch):
    fcn_model.train()          # train mode
    total_loss = 0.
    for batch_idx, (imgs, labels) in enumerate(train_loader):
... ...

在train函数中每当枚举train_loader时会调用voc_loader.py中的__getitem__函数，其定义如下所示：

    def __getitem__(self, index):
        data_file = self.files[self.split][index]    # 数据
        # load image
        img_file = data_file['img']
        img = PIL.Image.open(img_file)
        img = np.array(img, dtype=np.uint8)
        # load label
        lbl_file = data_file['lbl']
        lbl = PIL.Image.open(lbl_file)
        lbl = np.array(lbl, dtype=np.uint8)

        lbl[lbl == 255] = 0
        # augment
        img, lbl = self.randomFlip(img, lbl)
        img, lbl = self.randomCrop(img, lbl)
        img, lbl = self.resize(img, lbl)

        if self._transform:
            return self.transform(img, lbl)
        else:
            return img, lbl

 该函数的第一行就是根据split值（train或val）来读取相应的数据。每个数据data_file包括一张样本图片jpg和一张标注结果图片png，其二进制像素值分别赋给了img和lbl。

紧接着对img和lbl做数据增强，包括randomFlip（左右镜像）， randomCrop（随机剪裁）以及resize（成统一分辨率）。

最后做了一个transform，其具体实现如下所示：

def transform(self, img, lbl):
        img = img[:, :, ::-1]          # RGB -> BGR
        img = img.astype(np.float64)
        img -= self.mean_bgr
        img = img.transpose(2, 0, 1)   # whc -> cwh
        img = torch.from_numpy(img).float()
        lbl = torch.from_numpy(lbl).long()
        return img, lbl

主要是将rgb通道顺序调转成bgr，而且每个像素值都要减去mean值，并作了一个转置，最后就是改变img和lbl的数值类型。经过这样的数据处理后，才可以输入到FCN模型中进行训练。