color.py

Go to the documentation of this file.

import numpy as np
 
 
# Taken from https://github.com/wkentaro/imgviz/blob/master/imgviz/label.py  # NOQA
def labelcolormap(N=256):
    """Label colormap.
    Parameters
    ----------
    N: int
        Number of labels (default: 256).
 
    Returns
    -------
    cmap: numpy.ndarray, (N, 3), numpy.uint8
        Label id to colormap.
    """
 
    def bitget(byteval, idx):
        shape = byteval.shape + (8,)
        return np.unpackbits(byteval).reshape(shape)[..., -1 - idx]
 
    i = np.arange(N, dtype=np.uint8)
    r = np.full_like(i, 0)
    g = np.full_like(i, 0)
    b = np.full_like(i, 0)
 
    i = np.repeat(i[:, None], 8, axis=1)
    i = np.right_shift(i, np.arange(0, 24, 3)).astype(np.uint8)
    j = np.arange(8)[::-1]
    r = np.bitwise_or.reduce(np.left_shift(bitget(i, 0), j), axis=1)
    g = np.bitwise_or.reduce(np.left_shift(bitget(i, 1), j), axis=1)
    b = np.bitwise_or.reduce(np.left_shift(bitget(i, 2), j), axis=1)
 
    cmap = np.stack((r, g, b), axis=1).astype(np.uint8)
    return cmap