Estoy tratando de convertir el mapa de color de un contorno generado desde una aplicación que no sea Python. Intenté usar el código de Matthias Bussonnier disponibleaqu, pero no puede darme una conversión completa. Traté de truncar el mapa de colores para obtener una conversión completa, pero nuevamente no me da una conversión completa.

MWE

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.colors as colors
import matplotlib.image as mpimg
from scipy.spatial import cKDTree
import matplotlib


import matplotlib.cm as mplcm

def truncate_colormap(cmap, minval=0.0, maxval=1.0, n=100):
    new_cmap = colors.LinearSegmentedColormap.from_list(
        'trunc({n},{a:.2f},{b:.2f})'.format(n=cmap.name, a=minval, b=maxval),
        cmap(np.linspace(minval, maxval, n)))
    return new_cmap
cmap = plt.get_cmap('jet')
cmap = truncate_colormap(cmap, 0.1, 0.9)
img = mpimg.imread('./test.png')[:,:,:3]

#@interact(sub=(0, 500), d=(0,1,0.05))
def convert(sub=256,d=0.1, cin=cmap, cout='viridis'):
    viridis = plt.get_cmap(cout)
    jet = plt.get_cmap(cin)
    jet256 = colors.makeMappingArray(sub, jet)[:, :3]
    K = cKDTree(jet256)
    oshape = img.shape
    img_data = img.reshape((-1,3))
    res = K.query(img_data, distance_upper_bound=d)
    indices = res[1]
    l = len(jet256)
    indices = indices.reshape(oshape[:2])
    remapped = indices

    indices.max()

    mask = (indices == l)

    remapped = remapped / (l-1)
    mask = np.stack( [mask]*3, axis=-1)

    blend = np.where(mask, img, viridis(remapped)[:,:,:3])
    fig, ax = plt.subplots()
    fig.set_figheight(10)
    fig.set_figwidth(10)
    ax.imshow(blend)
    fig.savefig('viridize.pdf')
convert()    

Input image

Output image

¿Cómo obtengo una conversión completa del mapa de colores (jet en este caso) a viridis con Python?