cam.py

"""Converts between the RGB and CIECAM02 color spaces."""

from collections import namedtuple
from functools import partial

import colour
from colour.utilities import tsplit, tstack
import numpy as np
from scipy.optimize import fmin_l_bfgs_b


Conditions = namedtuple('Conditions', 'Y_w L_A Y_b surround')
DARK_BG = Conditions(80, 16, 0.8, colour.CIECAM02_VIEWING_CONDITIONS['Average'])
NEUTRAL_BG = Conditions(80, 16, 16, colour.CIECAM02_VIEWING_CONDITIONS['Average'])
LIGHT_BG = Conditions(80, 16, 80, colour.CIECAM02_VIEWING_CONDITIONS['Average'])


def apow(x, power):
    """Raises x to the given power, treating negative numbers in a mirrored fashion."""
    return np.abs(x)**power * np.sign(x)


RGB_CS = colour.RGB_Colourspace('rgb',
                                colour.models.sRGB_COLOURSPACE.primaries,
                                colour.models.sRGB_COLOURSPACE.whitepoint,
                                encoding_cctf=partial(apow, power=1/2.2),
                                decoding_cctf=partial(apow, power=2.2),
                                use_derived_RGB_to_XYZ_matrix=True,
                                use_derived_XYZ_to_RGB_matrix=True)


def rgb_to_xyz(RGB):
    """Converts from the RGB colorspace to XYZ tristimulus values."""
    return colour.RGB_to_XYZ(RGB, RGB_CS.whitepoint, RGB_CS.whitepoint,
                             RGB_CS.RGB_to_XYZ_matrix, decoding_cctf=RGB_CS.decoding_cctf)


def xyz_to_rgb(XYZ):
    """Converts from XYZ tristimulus values to the RGB colorspace."""
    return colour.XYZ_to_RGB(XYZ, RGB_CS.whitepoint, RGB_CS.whitepoint,
                             RGB_CS.XYZ_to_RGB_matrix, encoding_cctf=RGB_CS.encoding_cctf)


def rgb_to_cam(RGB, conds=NEUTRAL_BG):
    """Converts from the RGB colorspace to CIECAM02."""
    XYZ = rgb_to_xyz(RGB)
    XYZ_w = rgb_to_xyz([1, 1, 1])
    Y_w, L_A, Y_b, surround = conds
    cam = colour.XYZ_to_CIECAM02(XYZ * Y_w, XYZ_w * Y_w, L_A, Y_b, surround,
                                 discount_illuminant=True)
    return tstack([cam.J, cam.C, cam.h])


def cam_to_rgb(JCh, conds=NEUTRAL_BG):
    """Converts from CIECAM02 to the RGB colorspace."""
    XYZ_w = rgb_to_xyz([1, 1, 1])
    Y_w, L_A, Y_b, surround = conds
    spec = colour.CIECAM02_Specification(*tsplit(JCh))
    XYZ = colour.CIECAM02_to_XYZ(spec, XYZ_w * Y_w, L_A, Y_b, surround,
                                 discount_illuminant=True) / Y_w
    return xyz_to_rgb(XYZ)


def jch_to_jab(JCh):
    """Converts from cylindrical (JCh) to Cartesian (Jab) coordinates."""
    J, C, h = tsplit(JCh)
    h_ = np.deg2rad(h)
    a = C * np.cos(h_)
    b = C * np.sin(h_)
    return tstack([J, a, b])


def jab_to_jch(Jab):
    """Converts from Cartesian (Jab) to cylindrical (JCh) coordinates."""
    J, a, b = tsplit(Jab)
    C = np.sqrt(a**2 + b**2)
    h = np.rad2deg(np.arctan2(b, a)) % 360
    return tstack([J, C, h])


def constrain_to_gamut(rgb, conds=NEUTRAL_BG):
    """Constrains the given RGB colorspace color to lie within the RGB colorspace gamut,
    minimizing the CIECAM02 distance between the input out-of-gamut color and the output in-gamut
    color."""
    x = np.clip(rgb, 0, 1)
    if (rgb == x).all():
        return x
    Jab = jch_to_jab(rgb_to_cam(rgb, conds))
    def loss(rgb_):
        Jab_ = jch_to_jab(rgb_to_cam(rgb_, conds))
        return sum((Jab - Jab_)**2)
    x_opt, _, _ = fmin_l_bfgs_b(loss, x, approx_grad=True, bounds=[(0, 1)]*3)
    return x_opt


def modify(rgb, conds_src=NEUTRAL_BG, conds_dst=NEUTRAL_BG,
           scale_lightness=1, scale_chroma=1, rotate_hue=0):
    """Modifies the given RGB colorspace color in various ways."""
    JCh = rgb_to_cam(rgb, conds_src)
    JCh[..., 0] *= scale_lightness
    JCh[..., 1] *= scale_chroma
    JCh[..., 2] += rotate_hue
    JCh[..., 2] %= 360
    rgb_ = cam_to_rgb(JCh, conds_dst)
    return np.apply_along_axis(partial(constrain_to_gamut, conds=conds_dst), -1, rgb_)