Demonstration of N overlays on a surface (standard GPU method)

pyvista_overlay_gpu_compositing.py

import numpy as np
import pyvista
from pyvista import examples
from matplotlib.colors import ListedColormap


# BEGIN: Inspired by Vispy
# https://github.com/vispy/vispy
_rgb2xyz_norm = np.array([[0.43395276, 0.212671, 0.01775791],
                         [0.37621941, 0.71516, 0.10947652],
                         [0.18982783, 0.072169, 0.87276557]])


def _check_color_dim(val):
    """Ensure val is Nx(n_col), usually Nx3"""
    val = np.atleast_2d(val)
    if val.shape[1] not in (3, 4):
        raise RuntimeError('Value must have second dimension of size 3 or 4')
    return val, val.shape[1]


def _rgb_to_lab(rgbs):
    rgbs, n_dim = _check_color_dim(rgbs)
    # convert RGB->XYZ
    xyz = rgbs[:, :3].copy()  # a misnomer for now but will end up being XYZ
    over = xyz > 0.04045
    xyz[over] = ((xyz[over] + 0.055) / 1.055) ** 2.4
    xyz[~over] /= 12.92
    xyz = np.dot(xyz, _rgb2xyz_norm)
    over = xyz > 0.008856
    xyz[over] = xyz[over] ** (1. / 3.)
    xyz[~over] = 7.787 * xyz[~over] + 0.13793103448275862

    # Convert XYZ->LAB
    L = (116. * xyz[:, 1]) - 16
    a = 500 * (xyz[:, 0] - xyz[:, 1])
    b = 200 * (xyz[:, 1] - xyz[:, 2])
    labs = [L, a, b]
    # Append alpha if necessary
    if n_dim == 4:
        labs.append(np.atleast1d(rgbs[:, 3]))
    labs = np.array(labs, order='F').T  # Becomes 'C' order b/c of .T
    return labs
# END: Inspired by Vispy


# From: https://github.com/vispy/vispy/issues/415#issuecomment-471491952
def distinguishable_colors(ncolors, bg=[0, 0, 0]):
    n_grid = 30
    x = np.linspace(0, 1, n_grid)
    R, G, B = np.meshgrid(x, x, x)

    rgb = np.column_stack((R.flatten(),
                           G.flatten(),
                           B.flatten())).astype(np.float32)
    lab = _rgb_to_lab(rgb)

    bglab = np.array(bg).astype(np.float32)
    bglab = _rgb_to_lab(bglab)

    mindist2 = np.full(rgb.shape[0], float('Inf'))
    for i in range(0, bglab.shape[0]):
        dX = lab - bglab[i, :]
        dist2 = np.sum(dX**2, axis=1)
        mindist2 = np.minimum(dist2, mindist2)

    colors = np.zeros((ncolors, 3))
    lastlab = bglab[-1, :]
    for i in range(0, ncolors):
        dX = lab - lastlab
        dist2 = np.sum(dX**2, axis=1)
        mindist2 = np.minimum(dist2, mindist2)
        index = mindist2.argmax()
        colors[i, :] = rgb[index, :]
        lastlab = lab[index, :]
    return colors


N = 10
bg = [1, 1, 1]
colors = distinguishable_colors(N, bg=bg)

surface = examples.download_cow().triangulate().subdivide(3)
points = surface.points
bounds = [points.min(), points.max()]
diff = (bounds[1] - bounds[0]) / N

# GPU compositing
plotter = pyvista.Plotter()

plotter.add_mesh(
    mesh=surface,
    color='white',
    smooth_shading=True,
)

kwargs = dict(
    rng=[0, 1],
    opacity=1.,
    rgba=False,
    smooth_shading=True,
)
for i in range(N):
    overlay = surface.copy()
    overlay_scalars = points[:, 0] < bounds[0] + (N - i) * diff
    overlay_scalars = overlay_scalars.astype(np.int)
    next_color = (*colors[i, :], 1.)
    overlay_cmap = ListedColormap(np.array([next_color, (1, 1, 1, 0)]))

    plotter.add_mesh(
        mesh=overlay,
        scalars=overlay_scalars,
        cmap=overlay_cmap,
        **kwargs
    )

plotter.show()