image_writer.py 文件源码

def map_to_colors(buff, cmap_name):
    try:
        lut = cmd.color_map_luts[cmap_name]
    except KeyError:
        try:
            # if cmap is tuple, then we're using palettable or brewer2mpl cmaps
            if isinstance(cmap_name, tuple):
                cmap = get_brewer_cmap(cmap_name)
            else:
                cmap = mcm.get_cmap(cmap_name)
            cmap(0.0)
            lut = cmap._lut.T
        except ValueError:
            raise KeyError(
                "Your color map (%s) was not found in either the extracted"
                " colormap file or matplotlib colormaps" % cmap_name)

    if isinstance(cmap_name, tuple):
        # If we are using the colorbrewer maps, don't interpolate
        shape = buff.shape
        # We add float_eps so that digitize doesn't go out of bounds
        x = np.mgrid[0.0:1.0+np.finfo(np.float32).eps:lut[0].shape[0]*1j]
        inds = np.digitize(buff.ravel(), x)
        inds.shape = (shape[0], shape[1])
        mapped = np.dstack([(v[inds]*255).astype('uint8') for v in lut])
        del inds
    else:
        x = np.mgrid[0.0:1.0:lut[0].shape[0]*1j]
        mapped = np.dstack(
                [(np.interp(buff, x, v)*255).astype('uint8') for v in lut ])
    return mapped.copy("C")