<reponame>KruskalLin/Autocalibration<filename>color_utils.py
from __future__ import division
import numpy as np
from matplotlib import cm
from matplotlib.colors import ListedColormap, LinearSegmentedColormap


def high_res_colormap(low_res_cmap, resolution=1000, max_value=1):
    # Construct the list colormap, with interpolated values for higer resolution
    # For a linear segmented colormap, you can just specify the number of point in
    # cm.get_cmap(name, lutsize) with the parameter lutsize
    x = np.linspace(0,1,low_res_cmap.N)
    low_res = low_res_cmap(x)
    new_x = np.linspace(0,max_value,resolution)
    high_res = np.stack([np.interp(new_x, x, low_res[:,i]) for i in range(low_res.shape[1])], axis=1)
    return ListedColormap(high_res)


def opencv_rainbow(resolution=1000):
    # Construct the opencv equivalent of Rainbow
    opencv_rainbow_data = (
        (0.000, (1.00, 0.00, 0.00)),
        (0.400, (1.00, 1.00, 0.00)),
        (0.600, (0.00, 1.00, 0.00)),
        (0.800, (0.00, 0.00, 1.00)),
        (1.000, (0.60, 0.00, 1.00))
    )

    return LinearSegmentedColormap.from_list('opencv_rainbow', opencv_rainbow_data, resolution)


COLORMAPS = {'rainbow': opencv_rainbow(),
             'magma': high_res_colormap(cm.get_cmap('magma').reversed()),
             'bone': cm.get_cmap('bone', 10000)}


def array2color(arr, max_value=None, colormap='rainbow'):
    if max_value is None:
        max_value = arr.max()

    if arr.shape[-1] == 1:
        arr = np.squeeze(arr, axis=2)

    if len(arr.shape) == 2:
        norm_array = arr / max_value
        array = COLORMAPS[colormap](norm_array).astype(np.float32)
        array = array.transpose(2, 0, 1)
    elif len(arr.shape) == 3:
        assert(arr.shape[-1] == 3)
        array = 0.5 + arr * 0.5
    return array

from __future__ import division
import numpy as np
from matplotlib import cm
from matplotlib.colors import ListedColormap, LinearSegmentedColormap


def high_res_colormap(low_res_cmap, resolution=1000, max_value=1):
    # Construct the list colormap, with interpolated values for higer resolution
    # For a linear segmented colormap, you can just specify the number of point in
    # cm.get_cmap(name, lutsize) with the parameter lutsize
    x = np.linspace(0,1,low_res_cmap.N)
    low_res = low_res_cmap(x)
    new_x = np.linspace(0,max_value,resolution)
    high_res = np.stack([np.interp(new_x, x, low_res[:,i]) for i in range(low_res.shape[1])], axis=1)
    return ListedColormap(high_res)


def opencv_rainbow(resolution=1000):
    # Construct the opencv equivalent of Rainbow
    opencv_rainbow_data = (
        (0.000, (1.00, 0.00, 0.00)),
        (0.400, (1.00, 1.00, 0.00)),
        (0.600, (0.00, 1.00, 0.00)),
        (0.800, (0.00, 0.00, 1.00)),
        (1.000, (0.60, 0.00, 1.00))
    )

    return LinearSegmentedColormap.from_list('opencv_rainbow', opencv_rainbow_data, resolution)


COLORMAPS = {'rainbow': opencv_rainbow(),
             'magma': high_res_colormap(cm.get_cmap('magma').reversed()),
             'bone': cm.get_cmap('bone', 10000)}


def array2color(arr, max_value=None, colormap='rainbow'):
    if max_value is None:
        max_value = arr.max()

    if arr.shape[-1] == 1:
        arr = np.squeeze(arr, axis=2)

    if len(arr.shape) == 2:
        norm_array = arr / max_value
        array = COLORMAPS[colormap](norm_array).astype(np.float32)
        array = array.transpose(2, 0, 1)
    elif len(arr.shape) == 3:
        assert(arr.shape[-1] == 3)
        array = 0.5 + arr * 0.5
    return array
# Construct the list colormap, with interpolated values for higer resolution # For a linear segmented colormap, you can just specify the number of point in # cm.get_cmap(name, lutsize) with the parameter lutsize # Construct the opencv equivalent of Rainbow