Plot diverse labels for segmentation image

diverse_cmap.py

def color_map(N=256, normalized=True):
    def bitget(byteval, idx):
        return (byteval & (1 << idx)) != 0

    dtype = 'float32' if normalized else 'uint8'
    cmap = np.zeros((N, 3), dtype=dtype)
    for i in range(N):
        r = g = b = 0
        c = i
        for j in range(8):
            r = r | (bitget(c, 0) << 7 - j)
            g = g | (bitget(c, 1) << 7 - j)
            b = b | (bitget(c, 2) << 7 - j)
            c = c >> 3

        cmap[i] = np.array([r, g, b])

    cmap = cmap / 255 if normalized else cmap
    return ListedColormap(cmap)

num_cats = 4
cmap = color_map(num_cats)
cats = ['cat']  * num_cats #List of categories
def make_plot(orig_im, res):
    global cats, cmap
    fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(20, 10))
    im_plt = ax[0].imshow(np.array(orig_im))
    cax = ax[1].imshow(res, cmap=cmap, vmax=num_cats, vmin=0)
    # Add colorbar, make sure to specify tick locations to match desired ticklabels
    cbar = fig.colorbar(cax, ticks=np.linspace(0.5, num_cats-0.5, num_cats))
    cbar.ax.set_yticklabels(cats)  # vertically oriented colorbar
    plt.show()