spec2rgb.py

"""
# --------------------------------------------
# spec2rgb: Convert spectral images to rgb images   
# --------------------------------------------

Python Script to map spectral image to a RGB representation.

# Source Code: https://github.com/hdspgroup/spec2rgb
"""

import numpy as np


def g(x, alpha, mu, sigma1, sigma2):
    sigma = (x < mu)*sigma1 + (x >= mu)*sigma2
    return alpha*np.exp((x-mu)**2 / (-2*(sigma**2)))


def component_x(x): return g(x, 1.056, 5998, 379, 310) + \
    g(x, 0.362, 4420, 160, 267) + g(x, -0.065, 5011, 204, 262)


def component_y(x): return g(x, 0.821, 5688, 469, 405) + \
    g(x, 0.286, 5309, 163, 311)


def component_z(x): return g(x, 1.217, 4370, 118, 360) + \
    g(x, 0.681, 4590, 260, 138)


def xyz_from_xy(x, y):
    """Return the vector (x, y, 1-x-y)."""
    return np.array((x, y, 1-x-y))


ILUMINANT = {
    'D65': xyz_from_xy(0.3127, 0.3291),
    'E':  xyz_from_xy(1/3, 1/3),
}

COLOR_SPACE = {
    'sRGB': (xyz_from_xy(0.64, 0.33),
             xyz_from_xy(0.30, 0.60),
             xyz_from_xy(0.15, 0.06),
             ILUMINANT['D65']),

    'fakeRGB': (xyz_from_xy(0.64, 0.33),
            xyz_from_xy(0.30, 0.60),
            xyz_from_xy(0.15, 0.06),
            ILUMINANT['D65']),

    'AdobeRGB': (xyz_from_xy(0.64, 0.33),
                 xyz_from_xy(0.21, 0.71),
                 xyz_from_xy(0.15, 0.06),
                 ILUMINANT['D65']),

    'AppleRGB': (xyz_from_xy(0.625, 0.34),
                 xyz_from_xy(0.28, 0.595),
                 xyz_from_xy(0.155, 0.07),
                 ILUMINANT['D65']),

    'UHDTV': (xyz_from_xy(0.708, 0.292),
              xyz_from_xy(0.170, 0.797),
              xyz_from_xy(0.131, 0.046),
              ILUMINANT['D65']),

    'CIERGB': (xyz_from_xy(0.7347, 0.2653),
               xyz_from_xy(0.2738, 0.7174),
               xyz_from_xy(0.1666, 0.0089),
               ILUMINANT['E']),
}


class ColourSystem:
    


    def __init__(self, start=380, end=750, num=100, cs='sRGB', custom_cs=None):

        # Chromaticities
        bands = np.linspace(start=start, stop=end, num=num)*10

        self.cmf = np.array([component_x(bands),
                             component_y(bands),
                             component_z(bands)])

        if custom_cs:
            self.red, self.green, self.blue, self.white = custom_cs
        else:
            self.red, self.green, self.blue, self.white = COLOR_SPACE[cs]

        # The chromaticity matrix (rgb -> xyz) and its inverse
        self.M = np.vstack((self.red, self.green, self.blue)).T
        self.MI = np.linalg.inv(self.M)

        # White scaling array
        self.wscale = self.MI.dot(self.white)

        # xyz -> rgb transformation matrix
        self.A  = self.MI / self.wscale[:, np.newaxis]
        self.GS = np.array([
            [0.3, 0.59, 0.11]
        ]).T

  
    def get_transform_matrix(self):

        XYZ = self.cmf
        RGB = XYZ.T @ self.A.T
        RGB = RGB / np.sum(RGB, axis=0, keepdims=True)
        return RGB


    def spec_to_rgb(self, spec):
        """Convert a spectrum to an rgb value."""
        M = self.get_transform_matrix()
        rgb = np.dot( spec , M)
        return rgb

    
    def spec_to_gs(self, spec):
        """Convert a spectrum to a greyscale value."""
       
        M = self.get_transform_matrix() 
        M = np.dot( M, self.GS )
        rgb = np.dot( spec, M )
        return rgb