import cv2
import numpy as np
from matplotlib import pyplot as plt


def apply_custom_colormap(image_gray, cmap=plt.get_cmap('seismic')):
    assert image_gray.dtype == np.uint8, 'must be np.uint8 image'
    if image_gray.ndim == 3: image_gray = image_gray.squeeze(-1)

    # Initialize the matplotlib color map
    sm = plt.cm.ScalarMappable(cmap=cmap)

    # Obtain linear color range
    color_range = sm.to_rgba(np.linspace(0, 1, 256), bytes=True)[:,2::-1]
    color_range = color_range.reshape(256,1,3)

    return cv2.applyColorMap(image_gray, color_range)

image_gray = cv2.imread('cage.png', cv2.IMREAD_GRAYSCALE)
image_bgr = apply_custom_colormap(image_gray, cmap=plt.get_cmap('bwr'))

cv2.imshow('image with colormap', image_bgr)
cv2.waitKey(0)