How can I change colors in contours (obtained from

I am trying to convert the color map of a contour generated from non-Python application. I tried using Matthias Bussonnier's code available here, but is unable to give me a full conversion. I tried to truncate the color map to give me a full conversion, but again does not give me a complete conversion.

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

How do I get a complete conversion of the color map (jet in this case) to viridis with Python?

As commented, the solution from How I can specify how rainbow color scheme should be converted to grayscale will work, but with some small modifications.

I.e. you need to apply your target colormap to the values optained from that solution and hence modify the resulting array size to be 3D.

The conditions for this to work are:

You know the colormap that the original image has been produced with (origin_cmap)
All colors in that image are either grey scale (axes, text etc.) or part of that origin_cmap. I.e. there should not be any other line plot or similar in addition in the figure.
The original colormap is unambiguous, i.e. does not contain the same color twice.
The full range of the original colormap has been used to create the input image and the full range of the target colormap will be aimed for. (This condition can be weakend though if needed, by specifying a different norm and/or range)

The following will hence "viridify" a given image.

import numpy as np
import matplotlib.colors
import matplotlib.pyplot as plt

image = plt.imread("https://i.stack.imgur.com/NyLq2.png")

def changecolormap(image, origin_cmap, target_cmap):
    r = np.linspace(0,1, 256)
    norm = matplotlib.colors.Normalize(0,1)
    mapvals = origin_cmap(norm(r))[:,:3]

    def get_value_from_cm(color):
        color=matplotlib.colors.to_rgb(color)
        #if color is already gray scale, dont change it
        if np.std(color) < 0.1:
            return color
        #otherwise return value from colormap
        distance = np.sum((mapvals - color)**2, axis=1)
        return target_cmap(r[np.argmin(distance)])[:3]

    newim = np.zeros_like(image)
    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            c = image[i,j,:3]
            newim[i,j, :3] =  get_value_from_cm(c)
    return newim


fig, (ax,ax2) = plt.subplots(ncols=2)
ax.imshow(image)
ax2.imshow(changecolormap(image, plt.cm.jet, plt.cm.viridis))

ax.axis("off")
ax2.axis("off")
plt.show()