tedit/shaders/white_balance.frag

#version 330 core

out vec4 FragColor;
in vec2 TexCoord;

uniform sampler2D InputTexture;
uniform float temperature; // Target Correlated Color Temperature (Kelvin, e.g., 1000-20000)
uniform float tint;        // Green-Magenta shift (-100 to +100)

// --- Constants ---

// sRGB Primaries to XYZ (D65)
const mat3 M_RGB_2_XYZ_D65 = mat3(
    0.4124564, 0.3575761, 0.1804375,
    0.2126729, 0.7151522, 0.0721750,
    0.0193339, 0.1191920, 0.9503041
);

// XYZ (D65) to sRGB Primaries
const mat3 M_XYZ_2_RGB_D65 = mat3(
     3.2404542, -1.5371385, -0.4985314,
    -0.9692660,  1.8760108,  0.0415560,
     0.0556434, -0.2040259,  1.0572252
);

// Standard Illuminant D65 XYZ (Normalized Y=1) - Our Target White
const vec3 WHITEPOINT_D65 = vec3(0.95047, 1.00000, 1.08883);

// Bradford CAT Matrices
const mat3 M_BRADFORD = mat3(
    0.8951,  0.2664, -0.1614,
   -0.7502,  1.7135,  0.0367,
    0.0389, -0.0685,  1.0296
);
const mat3 M_BRADFORD_INV = mat3(
    0.9869929, -0.1470543, 0.1599627,
    0.4323053,  0.5183603, 0.0492912,
   -0.0085287,  0.0400428, 0.9684867
);

// --- Helper Functions ---

// Approximate Kelvin Temperature to XYZ Chromaticity (xy) -> XYZ coordinates
// Uses simplified polynomial fits for different temperature ranges.
// Note: More accurate methods often use look-up tables or spectral calculations.
vec3 kelvinToXYZ(float kelvin) {
    float temp = clamp(kelvin, 1000.0, 20000.0);
    float x, y;

    // Calculate xy chromaticity coordinates based on temperature
    // Formulas from: http://www.brucelindbloom.com/index.html?Eqn_T_to_xy.html (with slight adaptations)
    float t = temp;
    float t2 = t * t;
    float t3 = t2 * t;

    // Calculate x coordinate
    if (t >= 1000.0 && t <= 4000.0) {
        x = -0.2661239e9 / t3 - 0.2343589e6 / t2 + 0.8776956e3 / t + 0.179910;
    } else { // t > 4000.0 && t <= 25000.0
        x = -3.0258469e9 / t3 + 2.1070379e6 / t2 + 0.2226347e3 / t + 0.240390;
    }

    // Calculate y coordinate based on x
    float x2 = x * x;
    float x3 = x2 * x;
    if (t >= 1000.0 && t <= 2222.0) {
         y = -1.1063814 * x3 - 1.34811020 * x2 + 2.18555832 * x - 0.18709;
    } else if (t > 2222.0 && t <= 4000.0) {
         y = -0.9549476 * x3 - 1.37418593 * x2 + 2.09137015 * x - 0.16748867;
    } else { // t > 4000.0 && t <= 25000.0
         y = 3.0817580 * x3 - 5.8733867 * x2 + 3.75112997 * x - 0.37001483;
    }

    // Convert xyY (Y=1) to XYZ
    if (y < 1e-6) return vec3(0.0); // Avoid division by zero
    float Y = 1.0;
    float X = (x / y) * Y;
    float Z = ((1.0 - x - y) / y) * Y;

    return vec3(X, Y, Z);
}

// Apply Bradford Chromatic Adaptation Transform
vec3 adaptXYZ(vec3 xyzColor, vec3 sourceWhiteXYZ, vec3 destWhiteXYZ) {
    vec3 sourceCone = M_BRADFORD * sourceWhiteXYZ;
    vec3 destCone = M_BRADFORD * destWhiteXYZ;

    // Avoid division by zero if source cone response is zero
    // (shouldn't happen with typical white points but good practice)
    if (sourceCone.r < 1e-6 || sourceCone.g < 1e-6 || sourceCone.b < 1e-6) {
        return xyzColor; // Return original color if adaptation is impossible
    }

    vec3 ratio = destCone / sourceCone;

    mat3 adaptationMatrix = M_BRADFORD_INV * mat3(ratio.x, 0, 0, 0, ratio.y, 0, 0, 0, ratio.z) * M_BRADFORD;

    return adaptationMatrix * xyzColor;
}

// --- Main White Balance Function ---
vec3 applyWhiteBalance(vec3 linearSRGBColor, float tempK, float tintVal) {

    // 1. Convert input Linear sRGB (D65) to XYZ D65
    vec3 inputXYZ = M_RGB_2_XYZ_D65 * linearSRGBColor;

    // 2. Calculate the XYZ white point of the source illuminant (from Kelvin temp)
    // This is the white we want to adapt *FROM*
    vec3 sourceWhiteXYZ = kelvinToXYZ(tempK);

    // 3. Apply Bradford CAT to adapt from sourceWhiteXYZ to D65
    vec3 adaptedXYZ = adaptXYZ(inputXYZ, sourceWhiteXYZ, WHITEPOINT_D65);

    // 4. Convert adapted XYZ (now relative to D65) back to Linear sRGB
    vec3 adaptedLinearSRGB = M_XYZ_2_RGB_D65 * adaptedXYZ;

    // 5. Apply Tint adjustment (Post-CAT approximation)
    // This shifts color balance along a Green<->Magenta axis.
    // We scale RGB components slightly based on the tint value.
    // A common method is to affect Green opposite to Red/Blue.
    if (abs(tintVal) > 0.01) {
        // Map tint -100..100 to a scaling factor, e.g., -0.1..0.1
        float tintFactor = tintVal / 1000.0; // Smaller scale factor for subtle tint

        // Apply scaling: Increase G for negative tint, Decrease G for positive tint
        // Compensate R and B slightly in the opposite direction.
        // Coefficients below are heuristic and may need tuning for perceptual feel.
        float rScale = 1.0 + tintFactor * 0.5;
        float gScale = 1.0 - tintFactor * 1.0;
        float bScale = 1.0 + tintFactor * 0.5;

        vec3 tintScaleVec = vec3(rScale, gScale, bScale);

        // Optional: Normalize scale vector to preserve luminance (roughly)
        // Calculate luminance of the scale vector itself
        float scaleLum = dot(tintScaleVec, vec3(0.2126, 0.7152, 0.0722));
        if (scaleLum > 1e-5) {
           tintScaleVec /= scaleLum; // Normalize
        }

        adaptedLinearSRGB *= tintScaleVec;
    }

    return adaptedLinearSRGB;
}

void main() {
    vec4 texColor = texture(InputTexture, TexCoord);
    vec3 linearInputColor = texColor.rgb; // Assuming input texture is already linear sRGB

    // Calculate white balanced color
    vec3 whiteBalancedColor = applyWhiteBalance(linearInputColor, temperature, tint);

    // Ensure output is non-negative
    whiteBalancedColor = max(whiteBalancedColor, vec3(0.0));

    FragColor = vec4(whiteBalancedColor, texColor.a);
}