前言
無論是哪個語義分割算法模型，比如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模型中進行訓練。