replace VOC helper with more general visualization utility

Colorize and visualize outputs by making images instead of relying on
embedded palettes. The palettes are now made in the PASCAL VOC style.

Note that other datasets can be handled by supplying a palette
as a no. of classes x 3 array.

vis.py

+import numpy as np
+
+def make_palette(num_classes):
+    """
+    Maps classes to colors in the style of PASCAL VOC.
+    Close values are mapped to far colors for segmentation visualization.
+    See http://host.robots.ox.ac.uk/pascal/VOC/voc2012/index.html#devkit
+
+    Takes:
+        num_classes: the number of classes
+    Gives:
+        palette: the colormap as a k x 3 array of RGB colors
+    """
+    palette = np.zeros((num_classes, 3), dtype=np.uint8)
+    for k in xrange(0, num_classes):
+        label = k
+        i = 0
+        while label:
+            palette[k, 0] |= (((label >> 0) & 1) << (7 - i))
+            palette[k, 1] |= (((label >> 1) & 1) << (7 - i))
+            palette[k, 2] |= (((label >> 2) & 1) << (7 - i))
+            label >>= 3
+            i += 1
+    return palette
+
+def color_seg(seg, palette):
+    """
+    Replace classes with their colors.
+
+    Takes:
+        seg: H x W segmentation image of class IDs
+    Gives:
+        H x W x 3 image of class colors
+    """
+    return palette[seg.flat].reshape(seg.shape + (3,))
+
+def vis_seg(img, seg, palette, alpha=0.5):
+    """
+    Visualize segmentation as an overlay on the image.
+
+    Takes:
+        img: H x W x 3 image in [0, 255]
+        seg: H x W segmentation image of class IDs
+        palette: K x 3 colormap for all classes
+        alpha: opacity of the segmentation in [0, 1]
+    Gives:
+        H x W x 3 image with overlaid segmentation
+    """
+    vis = np.array(img, dtype=np.float32)
+    mask = seg > 0
+    vis[mask] *= 1. - alpha
+    vis[mask] += alpha * palette[seg[mask].flat]
+    vis = vis.astype(np.uint8)
+    return vis

voc_helper.py

-import os
-import copy
-import glob
-import numpy as np
-
-from PIL import Image
-
-
-class voc:
-    def __init__(self, data_path):
-        # data_path is /path/to/PASCAL/VOC2011
-        self.dir = data_path
-        self.classes = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
-                        'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
-                        'diningtable', 'dog', 'horse', 'motorbike', 'person',
-                        'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']
-        # for paletting
-        reference_idx = '2008_000666'
-        palette_im = Image.open('{}/SegmentationClass/{}.png'.format(
-            self.dir, reference_idx))
-        self.palette = palette_im.palette
-
-    def load_image(self, idx):
-        im = Image.open('{}/JPEGImages/{}.jpg'.format(self.dir, idx))
-        return im
-
-    def load_label(self, idx):
-        """
-        Load label image as 1 x height x width integer array of label indices.
-        The leading singleton dimension is required by the loss.
-        """
-        label = Image.open('{}/SegmentationClass/{}.png'.format(self.dir, idx))
-        label = np.array(label, dtype=np.uint8)
-        label = label[np.newaxis, ...]
-        return label
-
-    def palette(self, label_im):
-        '''
-        Transfer the VOC color palette to an output mask for visualization.
-        '''
-        if label_im.ndim == 3:
-            label_im = label_im[0]
-        label = Image.fromarray(label_im, mode='P')
-        label.palette = copy.copy(self.palette)
-        return label