perfect-cmaps 1.0.0__py3-none-any.whl

perfect_cmaps/color_utils.py

@@ -0,0 +1,377 @@
+import matplotlib.colors as mcolors
+import numpy as np
+from scipy.stats import beta
+from scipy.interpolate import interp1d
+from typing import Tuple, Union
+from colour import sRGB_to_XYZ, XYZ_to_Lab, Lab_to_XYZ, XYZ_to_sRGB
+from matplotlib import pyplot as plt
+from matplotlib.colors import LinearSegmentedColormap
+
+from perfect_cmaps.optimization import genetic_algorithm
+from perfect_cmaps.storage import get_test_img_path
+
+
+RGB_WEIGHT = np.array([0.2989, 0.5870, 0.1140])
+SUPPORTED_L_PROFILES = [
+    "linear", 
+    "diverging",
+    "diverging_inverted",
+    "diverging_sharper",
+    "diverging_inverted_sharper",
+    "flat"
+]
+
+
+def diverging_envelope(x: np.ndarray, c: float = 4, x1: float = 0.25) -> np.ndarray:
+    result = np.zeros(x.shape)
+    m = x1 * (c - 2) / (1 - 2 * x1 + 1e-5)
+
+    below_x1 = x <= x1
+    result[below_x1] = c*x[below_x1]
+
+    before_middle = np.logical_and(x > x1, x <= 0.5)
+    result[before_middle] = m + 2*(1-m) * x[before_middle]
+
+    after_middle = np.logical_and(x > 0.5, x <= 1-x1)
+    result[after_middle] = 2 - m + 2*(m-1) * x[after_middle]
+
+    in_middle = np.logical_and(x > x1, x <= 1 - x1)
+    result[in_middle] = result.max()
+
+    after_x2 = x > 1-x1
+    result[after_x2] = c * (1 - x[after_x2])
+    return result
+
+
+def dying_cos(x: np.ndarray, offset: float = 0.0):
+    result = np.zeros_like(x)
+    values = np.mod(np.round(x - offset), 2) == 0
+    result[values] = (1 + np.cos(2*np.pi*(x[values]-offset)))/2
+    return result
+
+
+def unwrap_channels(ijk: Union[Tuple[int, int, int], None] = None) -> tuple:
+    if ijk is None:
+        ijk = (0, 1, 2)
+    
+    return ijk
+
+
+def linearize_rgba(rgba_colors: np.ndarray) -> np.ndarray:
+    lab_colors = XYZ_to_Lab(sRGB_to_XYZ(rgba_colors[:, :3]))
+    lab_colors[:, 0] = np.linspace(0, 100, lab_colors.shape[0])
+    rgb_colors_renormalized = Lab_to_sRGB(lab_colors)
+    rgb_colors_renormalized = np.concatenate(
+        (rgb_colors_renormalized, rgba_colors[:, 3:4]), axis=1
+    )
+    return rgb_colors_renormalized
+
+
+def cold_blooded(x: np.ndarray, ijk: Union[Tuple[int, int, int], None] = None) -> np.ndarray:
+    i, j, k = unwrap_channels(ijk)
+
+    result = np.zeros([*x.shape, 4])
+    period = 2*np.pi
+    result[..., j] = x**2 - x/period * np.sin(period*x)
+    result[..., k] = ((1 - np.cos(period/2*x))/2)**2
+    result[..., i] = 3*x**2 - result[..., j] - result[..., k]
+    result[..., :3] = np.sqrt(result[..., :3])
+
+    result[..., :3] /= np.max(np.max(result))
+    result[..., 3] = 1
+    
+    return result
+
+
+def cold_blooded_l(x: np.ndarray, ijk: Union[Tuple[int, int, int], None] = None) -> np.ndarray:
+    normal_colors = cold_blooded(x, ijk)
+    return linearize_rgba(normal_colors)
+
+
+def copper_salt(x: np.ndarray, ijk: Union[Tuple[int, int, int], None] = None) -> np.ndarray:
+    i, j, k = unwrap_channels(ijk)
+
+    result = np.zeros([*x.shape, 4])
+    dist = beta(2, 4)
+    pdf = dist.pdf
+    envelope = diverging_envelope(x, c=4, x1=0.5)
+    result[..., i] = pdf(x)
+    result[..., k] = pdf(1-x)
+    result[..., i] /= np.max(result[..., i])
+    result[..., k] /= np.max(result[..., k])
+    result[..., j] = np.sqrt(3 * envelope - result[..., i] ** 2 - result[..., k] ** 2)
+    weight = 1 / np.sqrt(RGB_WEIGHT)
+    
+    result[..., :3] = np.einsum(
+        'i...j,k...j->i...j', 
+        result[..., :3],
+        weight.reshape(1, *weight.shape)
+    )
+    result[..., :3] /= np.max(result)
+    result[..., 3] = 1
+    return result
+
+
+def get_lightness_profile(n_values: int, profile: str = "linear") -> np.ndarray:
+    if profile == "linear":
+        L_values = np.linspace(0, 100, n_values)
+    elif profile.startswith("diverging"):
+        base_envelope = np.linspace(-10, 10, n_values)
+
+        if profile == "diverging":
+            L_values = 100 - base_envelope ** 2
+
+        elif profile == "diverging_sharper":
+            L_values = 100 - (base_envelope ** 2 + np.abs(base_envelope * 10)) / 2
+
+        elif profile == "diverging_inverted":
+            L_values = base_envelope ** 2
+        
+        elif profile == "diverging_inverted_sharper":
+            L_values = (base_envelope ** 2 + np.abs(base_envelope * 10)) / 2
+
+    elif profile == "flat":
+        # Use a flat lightness profile, e.g., L = 50
+        L_values = np.full(n_values, 50)
+    else:
+        # Default to linear
+        L_values = np.linspace(0, 100, n_values)
+    
+    return L_values
+
+
+def interpolate_lab(
+    control_points: np.ndarray, 
+    num_values: int = 1000, 
+    profile: str = "linear", 
+    interpolation: str = "quadratic"
+) -> np.ndarray:
+    # Interpolate between the control points
+    lab_colors = np.zeros((num_values, 3))
+    space = np.linspace(0, 1, control_points.shape[0])
+    for i in range(2):
+        interpolator = interp1d(space, control_points[:, i], kind=interpolation)
+        lab_colors[:, i + 1] = interpolator(np.linspace(0, 1, num_values))
+
+    lab_colors[:, 0] = get_lightness_profile(num_values, profile)
+    return lab_colors
+
+
+def rgb_to_grayscale(rgb: np.ndarray) -> np.ndarray:
+    """Convert RGB values to grayscale using luminosity values of RGB channels."""
+    return np.sqrt(np.dot(rgb[...,:3] ** 2, RGB_WEIGHT))
+
+
+def rgb_to_grayscale_lightness(rgb: np.ndarray) -> np.ndarray:
+    """
+    Convert RGB values to grayscale values representing perceived lightness (L*).
+
+    Parameters
+    ----------
+    rgb : ndarray
+        An array of RGB values in sRGB color space, with values in the range [0, 1].
+
+    Returns
+    -------
+    L_normalized : ndarray
+        Grayscale values representing the L* component, normalized to [0, 1].
+    """
+    # Ensure RGB values are within [0, 1]
+    rgb = np.clip(rgb, 0, 1)
+
+    # Convert sRGB to XYZ
+    XYZ = sRGB_to_XYZ(rgb)
+
+    # Convert XYZ to CIELAB
+    Lab = XYZ_to_Lab(XYZ)
+
+    # Extract the L* component
+    L = Lab[..., 0]  # L* ranges from 0 to 100
+
+    # Normalize L* to [0, 1]
+    L_normalized = L / 100.0
+
+    return L_normalized
+
+
+def truncate_colormap(
+        cmap: mcolors.LinearSegmentedColormap, 
+        minval: float = 0.0, 
+        maxval: float = 1.0, 
+        n: int = 1000
+    ) -> mcolors.LinearSegmentedColormap:
+
+    new_cmap = mcolors.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
+
+
+def Lab_to_sRGB(lab_colors: np.ndarray) -> np.ndarray:
+    """
+    Convert CIE-LAB colors to sRGB color space.
+    
+    Parameters
+    ----------
+    lab_colors : ndarray
+        An array of Lab colors with shape (..., 3).
+    
+    Returns
+    -------
+    rgb_colors : ndarray
+        An array of sRGB colors with values in [0, 1].
+    """
+    # Convert Lab to XYZ
+    XYZ = Lab_to_XYZ(lab_colors)
+    
+    # Convert XYZ to sRGB
+    RGB = XYZ_to_sRGB(XYZ)
+    
+    return RGB
+
+
+def find_valid_L_range(a_star: float, b_star: float) -> Tuple[float, float]:
+    # Vectorized L* samples
+    L_samples = np.linspace(0, 100, 100)
+    # Create an array of Lab colors
+    lab = np.column_stack((L_samples, np.full_like(L_samples, a_star), np.full_like(L_samples, b_star)))
+    
+    # Convert all Lab colors to RGB at once
+    rgb = Lab_to_sRGB(lab)
+    
+    # Determine which colors are in gamut
+    in_gamut = np.all((rgb >= 0) & (rgb <= 1), axis=1)
+    
+    if np.any(in_gamut):
+        L_in_gamut = L_samples[in_gamut]
+        L_min_valid = L_in_gamut.min()
+        L_max_valid = L_in_gamut.max()
+    else:
+        L_min_valid = None
+        L_max_valid = None
+    
+    return L_min_valid, L_max_valid
+
+
+def find_valid_L_range_coarse_optim(a_star: float, b_star: float) -> Tuple[float, float]:
+    # Coarse sampling
+    L_samples_coarse = np.linspace(0, 100, 50)
+    lab_coarse = np.column_stack((L_samples_coarse, np.full_like(L_samples_coarse, a_star), np.full_like(L_samples_coarse, b_star)))
+    rgb_coarse = Lab_to_sRGB(lab_coarse)
+    in_gamut_coarse = np.all((rgb_coarse >= 0) & (rgb_coarse <= 1), axis=1)
+    
+    if np.any(in_gamut_coarse):
+        # Approximate L_min and L_max
+        L_in_gamut = L_samples_coarse[in_gamut_coarse]
+        L_min_approx = L_in_gamut.min()
+        L_max_approx = L_in_gamut.max()
+        
+        # Fine sampling near L_min_approx
+        L_samples_min = np.linspace(max(L_min_approx - 1, 0), L_min_approx, 20)
+        lab_min = np.column_stack((L_samples_min, np.full_like(L_samples_min, a_star), np.full_like(L_samples_min, b_star)))
+        rgb_min = Lab_to_sRGB(lab_min)
+        in_gamut_min = np.all((rgb_min >= 0) & (rgb_min <= 1), axis=1)
+        L_min_valid = L_samples_min[in_gamut_min].min() if np.any(in_gamut_min) else L_min_approx
+        
+        # Fine sampling near L_max_approx
+        L_samples_max = np.linspace(L_max_approx, min(L_max_approx + 1, 100), 20)
+        lab_max = np.column_stack((L_samples_max, np.full_like(L_samples_max, a_star), np.full_like(L_samples_max, b_star)))
+        rgb_max = Lab_to_sRGB(lab_max)
+        in_gamut_max = np.all((rgb_max >= 0) & (rgb_max <= 1), axis=1)
+        L_max_valid = L_samples_max[in_gamut_max].max() if np.any(in_gamut_max) else L_max_approx
+    else:
+        L_min_valid = None
+        L_max_valid = None
+    
+    return L_min_valid, L_max_valid
+
+
+def rgb_renormalized_lightness(
+    lab_colors: np.ndarray, 
+    population_size: int = 100, 
+    num_generations: int = 200,
+    verbose: bool = False,
+) -> np.ndarray:
+    
+    # Compute L_min and L_max for each color
+    n_colors = lab_colors.shape[0]
+    L_min = np.zeros(n_colors)
+    L_max = L_min.copy()
+
+    for i in range(n_colors):
+        a_star = lab_colors[i, 1]
+        b_star = lab_colors[i, 2]
+        # Compute L_min[i] and L_max[i] for this (a*, b*) pair
+        L_min[i], L_max[i] = find_valid_L_range(a_star, b_star)
+
+    L_intended = lab_colors[:, 0]
+
+    # Optimize m and c
+    gene_limits = np.array([[0.0, 100.0], [-100.0, 100.0]])
+    genes, best_fitness = genetic_algorithm(population_size, num_generations, gene_limits, L_intended, L_min, L_max)
+
+    if best_fitness == 0.0:
+        print("Optimization failed for chosen lightness profile.")
+        print("Colormap is not perfectly perceptually uniform.")
+        return Lab_to_sRGB(lab_colors), 1.0, 0.0
+    
+    if verbose:
+        print(f"Optimized colormap lightness range: {100 * best_fitness:.2f} %")
+
+    L_adjusted = L_intended * genes[0] + genes[1]
+
+    # Update L* values in lab_colors
+    lab_colors[:, 0] = L_adjusted
+
+    # Convert adjusted Lab colors to RGB
+    return Lab_to_sRGB(lab_colors)
+
+
+def plot_colormap(colormap: LinearSegmentedColormap, num_points: int = 1000):
+    gradient = np.linspace(0, 1, num_points)
+    gradient = np.vstack((gradient, gradient))
+
+    # Get RGB values from colormap
+    gradient_rgb = colormap(gradient)
+
+    # Convert the RGB values to grayscale
+    lightness = rgb_to_grayscale_lightness(gradient_rgb[0][:, :3])
+    luminance = rgb_to_grayscale(gradient_rgb[0][:, :3])
+    gradient_gray = np.vstack((lightness, lightness))
+
+    # Plotting
+    fig = plt.figure(figsize=(20, 10), constrained_layout=True)
+    gs = fig.add_gridspec(3, 2)
+
+    # Define the axes
+    ax0 = fig.add_subplot(gs[0, 0])  # Top-left subplot
+    ax1 = fig.add_subplot(gs[1, 0])  # Mid-left subplot
+    ax2 = fig.add_subplot(gs[2, 0])  # Bottom-left subplot
+    ax3 = fig.add_subplot(gs[:, 1])  # Right subplot spanning all rows
+
+    # Plot RGB gradient on ax0
+    ax0.imshow(gradient_rgb, aspect="auto", vmin=0.0, vmax=1.0)
+    ax0.set_title("RGB gradient", fontsize=20)
+    ax0.axis("off")
+
+    # Plot Grayscale gradient on ax1
+    ax1.imshow(gradient_gray, cmap="gray", aspect="auto", vmin=0.0, vmax=1.0)
+    ax1.set_title("Lightness gradient", fontsize=20)
+    ax1.axis("off")
+
+    test_image_path = get_test_img_path()
+    colormap_test_img_path = test_image_path / "colourmaptest.tif"
+    colormap_test_image = plt.imread(colormap_test_img_path)
+    ax2.imshow(colormap_test_image, cmap=colormap)
+    ax2.set_title("Colormap test image", fontsize=20)
+    ax2.axis("off")
+
+    # Plot the color channels and luminance on ax2
+    for c, color_string in zip(range(3), ["red", "green", "blue"]):
+        ax3.plot(gradient_rgb[0, :, c], color=color_string, label=color_string, linewidth=4)
+
+    ax3.plot(luminance, color="black", linestyle="-.", label="luminance")
+    ax3.plot(lightness, color="grey", label="lightness", linewidth=4)
+    ax3.set_title("Color channel intensities", fontsize=20)
+    ax3.legend(loc=0, fontsize=20)
+    plt.show()

perfect_cmaps/colormaps.py

@@ -0,0 +1,237 @@
+import matplotlib.colors as mcolors
+from matplotlib import pyplot as plt
+from matplotlib import colormaps
+
+import os, sys
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+import numpy as np
+from typing import Tuple, List
+import argparse
+import cv2
+from scipy.ndimage import gaussian_filter1d
+from pathlib import Path
+
+from perfect_cmaps.color_utils import *
+from perfect_cmaps.storage import load_data, get_test_img_path
+
+
+def get_colormap(
+        cmap_name: str, 
+        n: int = 100, 
+        ijk: Tuple = None, 
+        lightness: str = None,
+        interpolation: str = "quadratic",
+        smoothing: float = None
+    ) -> mcolors.LinearSegmentedColormap:
+    
+    """Function for getting custom colormaps. 
+    Two algorithmically generated colormaps are currently available:
+        - 'cold_blooded' a.k.a 'ectotherm'
+        - 'copper_salt'
+    The 'cold_blooded' map also has variants with completely linear lightness: 
+        - 'cold_blooded_l', or 'ectotherm_l'.
+        
+    Additionally, several control points in Lab format are also available in the "lab_control_points" folder.
+    These are custom generated through the script 'create_custom_cmap', and consist of a collection 
+    points from 'a' and 'b' channels in Lab, along with a lightness profile.
+    These points are used in an optimization task, which optimizes the envelope for maximal expressiveness.
+
+    You can change the lightness profile in this script by choosing a lightness profile,
+    but note that the optimization task might not be feasible for certain lightness profiles.
+
+    Args:
+        cmap_name (str): The name of the colormap, either from the control points folder,
+            or the two algorithmically generated colormaps.
+        n (int): Number of points in the colormap sequence.
+        ijk (Tuple, optional): Order of the channels for the two algorithmic colormaps.
+            Typically you would choose a permutation of (0, 1, 2), such as (2, 1, 0). 
+            Defaults to None, which amounts to the default permutation of the colormaps.
+        lightness (str, optional): Lightness profile - only applies for the Lab
+            control points. Defaults to None, which means the lightness profile 
+            the chosen colormap was saved with.
+        interpolation (str): Interpolation method for L values. Defaults to 'quadratic.'
+        smoothing (float): Gaussian smoothing of colormap in Lab space, 
+            applied after interpolation. Defaults to None.
+
+    Returns:
+        mcolors.LinearSegmentedColormap: Matplotlib-formatted colormap
+    """
+    space = np.linspace(0, 1, n)
+
+    if not cmap_name in CMAP_DICT.keys():
+        json_data = load_data(cmap_name)
+        control_points = np.array(json_data["points"])
+
+        if lightness is None:
+            lightness = json_data["lightness"]
+        else:
+            assert lightness in SUPPORTED_L_PROFILES, f"Lightness profile {lightness} not supported"
+
+        interpolated_values = interpolate_lab(control_points, n, lightness, interpolation)
+        if smoothing is not None and smoothing > 0.0:
+            sigma = smoothing * n
+            interpolated_values[:, 1:] = gaussian_filter1d(interpolated_values[:, 1:], sigma=sigma, axis=0)
+
+        rgb_values = rgb_renormalized_lightness(interpolated_values, 500, 500)
+    
+    else:
+        color_function = CMAP_DICT[cmap_name]
+        rgb_values = color_function(space, ijk)
+        reoptimize_map = False
+        
+        if color_function in (cold_blooded_l, cold_blooded) and lightness == "linear":
+            reoptimize_map = False
+        
+        elif lightness is not None:
+            reoptimize_map = True
+
+        if reoptimize_map:
+            lab_values = XYZ_to_Lab(sRGB_to_XYZ(rgb_values[:, :3]))
+            lab_values[:, 0] = get_lightness_profile(n, lightness)
+            rgb_values, _, _ = rgb_renormalized_lightness(lab_values, 200, 500)
+    
+    cdict = dict()
+
+    for num, col in enumerate(["red", "green", "blue"]):
+        col_list = [[space[i], rgb_values[i][num], rgb_values[i][num]] for i in range(n)]
+        cdict[col] = col_list
+    cmp = mcolors.LinearSegmentedColormap(cmap_name, segmentdata=cdict, N=n)
+    return cmp
+
+
+def plot_images_with_colormap(image_paths: List[Path], colormap='viridis'):
+    # Create a figure with two rows, one for greyscale and one for colormap images
+    num_images = len(image_paths)
+    
+    fig, axs = plt.subplots(2, num_images, figsize=(25, 10))
+    axs[0, 0].set_ylabel("Grayscale", fontsize=20)
+    axs[1, 0].set_ylabel("Custom colormap", fontsize=20)
+    
+    for idx, img_path in enumerate(image_paths):
+        # Load the image in greyscale
+        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
+        img = img[:int(img.shape[0] * 0.94)]
+        
+        # Display the greyscale image
+        axs[0, idx].imshow(img, cmap='gray', vmin=0, vmax=255)
+        
+        # Display the image with the specified colormap
+        axs[1, idx].imshow(img, cmap=colormap)
+        axs[0, idx].set_xticks([])
+        axs[0, idx].set_yticks([])
+        axs[1, idx].set_xticks([])
+        axs[1, idx].set_yticks([])
+    
+    plt.tight_layout()
+    plt.show()
+
+
+def parse_args():
+    arg_parser = argparse.ArgumentParser(add_help=False)
+    arg_parser.add_argument(
+        "--colormap", 
+        "-c", 
+        type=str, 
+        default="ectotherm_l", 
+        help="The name of the colormap to be loaded"
+    )
+    arg_parser.add_argument(
+        "--num_points", 
+        "-n", 
+        type=int, 
+        default=1000, 
+        help="Number of points in the colormap spectrum"
+    )
+    arg_parser.add_argument(
+        "--lightness", 
+        "-l", 
+        type=str, 
+        default=None, 
+        help="The lightness profile of the colormap. Allowed values are: " + 
+             "linear, diverging, diverging_inverted, flat."
+    )
+    arg_parser.add_argument(
+        "--interpolation",
+        "-i",
+        type=str,
+        default="quadratic",
+        help="Interpolation method for Lab control points"
+    )
+    arg_parser.add_argument(
+        "--ijk",
+        type=int,
+        nargs='+',
+        default=(0, 1, 2),
+        help="Channel permutation for algorithmically generated colormaps"
+    )
+    arg_parser.add_argument(
+        "--smoothing",
+        "-s",
+        type=float,
+        default=None,
+        help="Gaussian smoothing of colormap in Lab space, applied after interpolation"
+    )
+    arg_parser.add_argument(
+        "--mpl_cmap",
+        "-mpl",
+        action="store_true",
+        default=False,
+        help="If specified, looks for a matplotlib colormap instead of one made with this library"
+    )
+    arg_parser.add_argument(
+        "--help", 
+        "-h", 
+        action="help", 
+        default=argparse.SUPPRESS,
+        help="Python script for colormap handling with perceptually uniform colormaps. " +
+             "Run as a standalone script or import and use the function 'get_colormap' in your own code."
+    )
+    return arg_parser.parse_args()
+
+
+CMAP_DICT = {
+    "cold_blooded": cold_blooded,
+    "cold_blooded_l": cold_blooded_l,
+    "ectotherm": cold_blooded,
+    "ectotherm_l": cold_blooded_l,
+    "copper_salt": copper_salt
+}
+
+
+"""
+If used as a main script, you can try colormaps out on test images, 
+and inspect the RGB channel profile of the colormap along with lightness and luminance.
+"""
+def main():
+    args = parse_args()
+    
+    if args.mpl_cmap:
+        cmap = colormaps[args.colormap]
+    else:
+        cmap = get_colormap(
+            args.colormap, 
+            args.num_points, 
+            ijk=args.ijk,
+            lightness=args.lightness, 
+            interpolation=args.interpolation,
+            smoothing=args.smoothing
+        )
+
+    plot_colormap(cmap, args.num_points)
+
+    test_images = [
+        "charcoal4_49721153188_o-scaled.jpg",
+        "Lily-Pollen.png",
+        "melamine-foam_49722002717_o-scaled.jpg",
+        "PhenomPharos-Gallery_Pyrite.jpg",
+        "PhenomXLGallery_5.jpg",
+    ]
+    
+    test_image_path = get_test_img_path()
+    test_images = [test_image_path / "sem" / img_name for img_name in test_images]
+    plot_images_with_colormap(test_images, cmap)
+
+
+if __name__ == "__main__":
+    main()