Full repository refactor

Author: f1recracker

TODO

@@ -1,4 +1,4 @@
 
-from dataset.bdd import *
+from dataset.berkeley_deepdrive import *
 from dataset.utils import *
 from dataset.transforms import *

dataset/__init__.py

@@ -0,0 +1,84 @@
+''' Berkeley Deepdrive Segmentation Dataset loader '''
+
+import os
+import re
+
+from PIL import Image
+import torch
+from torch.utils.data import Dataset
+
+from dataset.utils import listdir
+
+class BDDSegmentationDataset(Dataset):
+    ''' Dataset loader for Berkeley Deepdrive Segmentation dataset '''
+
+    def __init__(self, path, split, transforms=None):
+        assert split in ['train', 'val', 'test'], 'split must be one of: {train, val, test}'
+        image_re = re.compile(r'(.*)\.jpg')
+        label_re = re.compile(r'(.*)_train_id\.png')
+        images = sorted(listdir(os.path.join(path, 'seg/images', split), image_re))
+        labels = sorted(listdir(os.path.join(path, 'seg/labels', split), label_re))
+        for (image, label) in zip(images, labels):
+            assert (image_re.match(os.path.basename(image)).group(1) ==
+                    label_re.match(os.path.basename(label)).group(1))
+        self.images, self.labels = images, labels
+        self.transforms = transforms
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, key):
+        image = Image.open(self.images[key])
+        label = Image.open(self.labels[key])
+        if self.transforms:
+            image, label = self.transforms(image, label)
+        return image, label
+
+
+def bdd_palette(labels):
+    ''' Applies a color palette to either a single label
+        tensor or a batch of tensors '''
+    assert len(labels.shape) in [2, 3], 'Invalid labels shape'
+
+    # pylint: disable=bad-whitespace
+    color_map = torch.Tensor([
+        [128,  67, 125], # Road
+        [247,  48, 227], # Sidewalk
+        [ 72,  72,  72], # Building
+        [101, 103, 153], # Wall
+        [190, 151, 152], # Fence
+        [152, 152, 152], # Pole
+        [254, 167,  56], # Light
+        [221, 217,  55], # Sign
+        [106, 140,  51], # Vegetation
+        [146, 250, 157], # Terrain
+        [ 65, 130, 176], # Sky
+        [224,  20,  64], # Person
+        [255,   0,  25], # Rider
+        [  0,  22, 138], # Car
+        [  0,  11,  70], # Truck
+        [  0,  63,  98], # Bus
+        [  0,  82,  99], # Train
+        [  0,  36, 224], # Motorcycle
+        [121,  17,  38], # Bicycle
+        [  0,   0,   0]  # Other
+    ]).to(labels.device) / 255.0
+
+    batched_input = True
+    if len(labels.shape) == 2:
+        batched_input = False
+        labels = torch.unsqueeze(labels, 0)
+
+    # Convert ignore index to label 20
+    labels = torch.clamp(labels, 0, 20 - 1).long()
+
+    n, h, w = labels.shape
+    labels_one_hot = torch.zeros(n, 20, h, w).to(labels.device)
+    labels_one_hot.scatter_(1, torch.unsqueeze(labels, 1), 1)
+
+    color_labels = torch.einsum('nlhw,lc->nchw', labels_one_hot, color_map)
+
+    if not batched_input:
+        color_labels = torch.squeeze(color_labels, 0)
+
+    return color_labels

dataset/bdd.py

@@ -11,8 +11,8 @@ def transforms(img, seg, size=(360, 640), augment=True, hflip_prob=0.5,
                five_crop_prob=0.5, five_crop_scale=0.6,
                rotate_prob=0.5, max_rotate=30.0,
                tensor_output=True,
-               normalize_mean=torch.Tensor([0.0, 0.0, 0.0]),
-               normalize_std=torch.Tensor([1.0, 1.0, 1.0]),
+               normalize_mean=torch.Tensor([0.3518, 0.3932, 0.4011]),
+               normalize_std=torch.Tensor([0.2363, 0.2494, 0.2611]),
                _ignore_index=255):
     ''' BDD transforms pipeline '''
 

dataset/berkeley_deepdrive.py

@@ -1,6 +1,14 @@
 
+import os
+import re
+
 import torch
 
+def listdir(path, filter_=re.compile(r'.*')):
+    ''' Enumerates full paths of files in a directory matching a filter '''
+    return [os.path.join(path, f) for f in os.listdir(path) if filter_.match(f)]
+
+
 def median_frequency_balance(dataset, num_classes=19, ignore_index=255, _eps=1e-5):
     '''
     For more details refer to Section 6.3.2 in
@@ -14,3 +22,13 @@ def median_frequency_balance(dataset, num_classes=19, ignore_index=255, _eps=1e-
             frequency[cid] += torch.sum(seg == cid)
     frequency /= torch.sum(frequency)
     return torch.median(frequency) / frequency
+
+
+def mean_std(dataset):
+    ''' Returns the channel means and standard deviations
+        of the images in the dataset '''
+    mean, std = 0.0, 0.0
+    for image, _ in dataset:
+        mean += image.mean(dim=(1, 2)) # CHW -> C
+        std += image.view((3, -1)).std(dim=1) ** 2
+    return mean / len(dataset), (std / len(dataset)) ** 0.5

dataset/transforms.py

@@ -1,2 +1,2 @@
 
-from metrics.metrics import *
+from metrics.metrics import mean_iou, pixel_accuracy

dataset/utils.py

@@ -1,22 +1,45 @@
 
+#pylint: disable=invalid-name
+
 import torch
 
-def mean_iou(y_pred, y, logits_dim=1, ignore_index=255, eps=1e-8):
+def mean_iou(y_pred, y, num_classes, ignore_index=255):
     ''' Evaluates mean IoU between prediction and ground truth '''
-    y_pred = torch.argmax(y_pred, dim=logits_dim)
-    classes = set(torch.unique(torch.cat((y_pred, y))))
-    classes.discard(ignore_index)
-    mask = (y != ignore_index)
+    ignore_mask = (y != ignore_index)
+    y_pred, y = y_pred[ignore_mask], y[ignore_mask]
 
-    miou = 0.0
-    for i in classes:
-        intersect = torch.sum((y_pred[mask] == i) & (y[mask] == i)).float()
-        union = torch.sum((y_pred[mask] == i) | (y[mask] == i)).float()
-        miou += (intersect + eps) / (union + eps)
-    return (miou + eps) / (len(classes) + eps)
+    conf_matrix = _confusion_matrix(y, y_pred, num_classes)
+    true_pos = torch.diag(conf_matrix)
+    false_pos = torch.sum(conf_matrix, dim=0) - true_pos
+    false_neg = torch.sum(conf_matrix, dim=1) - true_pos
+    tp_fp_fn = true_pos + false_pos + false_neg
 
-def pixel_accuracy(y_pred, y, logits_dim=1, ignore_index=255):
+    exist_class_mask = tp_fp_fn > 0
+    true_pos, tp_fp_fn = true_pos[exist_class_mask], tp_fp_fn[exist_class_mask]
+    return torch.mean(true_pos / tp_fp_fn)
+
+def pixel_accuracy(y_pred, y, num_classes, ignore_index=255):
     ''' Evaluates pixel accuracy between prediction and ground truth '''
-    y_pred = torch.argmax(y_pred, dim=logits_dim)
     mask = (y != ignore_index)
-    return torch.sum(y[mask] == y_pred[mask]).float() / torch.sum(mask).float()
+    y_pred, y = y_pred[mask], y[mask]
+
+    conf_matrix = _confusion_matrix(y, y_pred, num_classes)
+    return torch.sum(torch.diag(conf_matrix)) / torch.sum(conf_matrix)
+
+# Helper functions
+
+def _one_hot(labels, num_classes, class_dim=1):
+    ''' Converts a labels tensor (NHW) into a one-hot tensor (NLHW) '''
+    labels = torch.unsqueeze(labels, class_dim)
+    labels_one_hot = torch.zeros_like(labels).repeat(
+        [num_classes if d == class_dim else 1
+         for d in range(len(labels.shape))])
+    labels_one_hot.scatter_(class_dim, labels, 1)
+    return labels_one_hot
+
+def _confusion_matrix(y_pred, y, num_classes):
+    ''' Computes the confusion matrix between two predicitons '''
+    b_size = y_pred.shape[0]
+    y, y_pred = _one_hot(y, num_classes), _one_hot(y_pred, num_classes)
+    y, y_pred = y.reshape(b_size, num_classes, -1), y_pred.reshape(b_size, num_classes, -1)
+    return torch.einsum('iaj,ibj->ab', y.float(), y_pred.float())

eval.py

@@ -1,5 +1,4 @@
-# pylint: disable=W0221,C0414,C0103
-
+# pylint: disable=arguments-differ, too-many-arguments
 ''' DeepLab V3+ '''
 
 import torch
@@ -83,7 +82,7 @@ def forward(self, x_in):
         logits = nn_func.interpolate(logits, size=x_in.shape[2:4],
                                      mode='bilinear', align_corners=True)
         return logits
-    
+
     def _init_weights(self):
         ''' Initializes weights of the model.
             - Conv2d parameters initialized using Kaiming normal

metrics/__init__.py

@@ -1,4 +1,4 @@
-# pylint: disable=W0221,C0414
+# pylint: disable=arguments-differ, too-many-arguments
 
 ''' Extensions to standard torch.nn primitives '''