Rust port of the GT7 Tone Mapping

bake.rs

use crate::GT7ToneMapping;

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum CurveMode {
    Linear,
    LinearSrgb,
    Pq,
    Srgb,
}

impl CurveMode {
    fn apply(&self, rgb: [f32; 3]) -> [f32; 3] {
        match self {
            CurveMode::Linear => rgb,
            CurveMode::LinearSrgb => rec2020_to_srgb(rgb),
            CurveMode::Pq => rgb.map(crate::inverse_eotf_pq),
            CurveMode::Srgb => rec2020_to_srgb(rgb).map(|value| {
                if value <= 0.0031308 {
                    12.92 * value
                } else {
                    1.055 * value.powf(1.0 / 2.4) - 0.055
                }
            }),
        }
    }
}

pub fn bake(
    size: u8,
    max_value: f32,
    tonemapper: GT7ToneMapping,
    curve_mode: CurveMode,
) -> Vec<[f32; 3]> {
    (0..size)
        .flat_map(|r| (0..size).flat_map(move |g| (0..size).map(move |b| [r, g, b])))
        .map(|[r, g, b]| {
            let tonemapped = tonemapper
                .apply([r, g, b].map(|x| (x as f32 / (size as f32 - 1.0)).powi(4) * max_value));
            curve_mode.apply(tonemapped)
        })
        .collect()
}

/// Convert from linear Rec.2020 to linear sRGB
#[allow(clippy::excessive_precision)]
fn rec2020_to_srgb(rgb: [f32; 3]) -> [f32; 3] {
    // Combined matrix: Rec.2020 (linear) -> sRGB (linear)
    const M: [[f32; 3]; 3] = [
        [1.6604915_f32, -0.5876412_f32, -0.0728503_f32],
        [-0.1245504_f32, 1.1328955_f32, -0.0083451_f32],
        [0.0181502_f32, -0.1005790_f32, 1.0824288_f32],
    ];

    let mut rec = [0.0_f32; 3];
    for i in 0..3 {
        rec[i] = M[i][0] * rgb[0] + M[i][1] * rgb[1] + M[i][2] * rgb[2];
    }
    rec
}

#[cfg(test)]
mod tests {
    use crate::UcsMode;

    use super::*;
    use std::io::Write;

    #[test]
    fn test_bake() {
        let lut = bake(
            33,
            100.0,
            GT7ToneMapping::new(None, UcsMode::ICtCp),
            CurveMode::Linear,
        );
        let mut file = std::fs::File::create("test.txt").unwrap();
        for line in lut {
            writeln!(file, "{:.6} {:.6} {:.6}", line[0], line[1], line[2]).unwrap();
        }
    }
}

tonemapping.rs

//! Rust port of the GT7 Tone Mapping (from `tonemapping.cpp`).

#![allow(clippy::excessive_precision)]

// Gran Turismo luminance-scale conversion helpers.
// In Gran Turismo, 1.0 in the linear frame-buffer corresponds to REFERENCE_LUMINANCE cd/m^2.
const REFERENCE_LUMINANCE: f32 = 100.0; // cd/m^2 <-> 1.0f
/// The SDR paper-white used by GT (cd/m^2)
pub const GRAN_TURISMO_SDR_PAPER_WHITE: f32 = 250.0; // cd/m^2

#[inline]
fn frame_buffer_value_to_physical(v: f32) -> f32 {
    v * REFERENCE_LUMINANCE
}

#[inline]
fn physical_to_frame_buffer(v: f32) -> f32 {
    v / REFERENCE_LUMINANCE
}

#[inline]
fn smooth_step(x: f32, edge0: f32, edge1: f32) -> f32 {
    if x <= edge0 {
        0.0
    } else if x >= edge1 {
        1.0
    } else {
        let t = (x - edge0) / (edge1 - edge0);
        t * t * (3.0 - 2.0 * t)
    }
}

#[inline]
fn chroma_curve(x: f32, a: f32, b: f32) -> f32 {
    1.0 - smooth_step(x, a, b)
}

// ST-2084 (PQ) EOTF / inverse EOTF
#[inline]
fn eotf_st2084(n: f32, exponent_scale_factor: f32) -> f32 {
    let n = n.clamp(0.0, 1.0);
    // SMPTE ST 2084 constants
    let m1: f32 = 0.159_301_76; // (2610 / 4096) / 4
    let m2: f32 = 78.843_75 * exponent_scale_factor; // (2523 / 4096) * 128
    let c1: f32 = 0.835_937_5; // 3424 / 4096
    let c2: f32 = 18.851_562; // (2413 / 4096) * 32
    let c3: f32 = 18.687_5; // (2392 / 4096) * 32
    let pq_c: f32 = 10_000.0; // cd/m^2

    let np = n.powf(1.0 / m2);
    let mut l = np - c1;
    if l < 0.0 {
        l = 0.0;
    }
    l /= c2 - c3 * np;
    l = l.powf(1.0 / m1);
    physical_to_frame_buffer(l * pq_c)
}

#[inline]
pub(crate) fn inverse_eotf_st2084(v: f32, exponent_scale_factor: f32) -> f32 {
    let m1: f32 = 0.159_301_76;
    let m2: f32 = 78.843_75 * exponent_scale_factor;
    let c1: f32 = 0.835_937_5;
    let c2: f32 = 18.851_562;
    let c3: f32 = 18.687_5;
    let pq_c: f32 = 10_000.0;

    let physical = frame_buffer_value_to_physical(v);
    let y = physical / pq_c;
    let ym = y.powf(m1);
    ((c1 + c2 * ym).log2() - (1.0 + c3 * ym).log2())
        .mul_add(m2, 0.0)
        .exp2()
}

// Color space conversions
const JZAZBZ_EXPONENT_SCALE_FACTOR: f32 = 1.7; // scale factor for exponent

#[inline]
fn rgb_to_ictcp(rgb: [f32; 3]) -> [f32; 3] {
    // Input: linear Rec.2020
    let l = (rgb[0] * 1688.0 + rgb[1] * 2146.0 + rgb[2] * 262.0) / 4096.0;
    let m = (rgb[0] * 683.0 + rgb[1] * 2951.0 + rgb[2] * 462.0) / 4096.0;
    let s = (rgb[0] * 99.0 + rgb[1] * 309.0 + rgb[2] * 3688.0) / 4096.0;

    let l_pq = inverse_eotf_st2084(l, 1.0);
    let m_pq = inverse_eotf_st2084(m, 1.0);
    let s_pq = inverse_eotf_st2084(s, 1.0);

    let i = (2048.0 * l_pq + 2048.0 * m_pq) / 4096.0;
    let ct = (6610.0 * l_pq - 13613.0 * m_pq + 7003.0 * s_pq) / 4096.0;
    let cp = (17933.0 * l_pq - 17390.0 * m_pq - 543.0 * s_pq) / 4096.0;
    [i, ct, cp]
}

#[inline]
fn ictcp_to_rgb(ictcp: [f32; 3]) -> [f32; 3] {
    let l = ictcp[0] + 0.008_609_04 * ictcp[1] + 0.111_03 * ictcp[2];
    let m = ictcp[0] - 0.008_609_04 * ictcp[1] - 0.111_03 * ictcp[2];
    let s = ictcp[0] + 0.560_031 * ictcp[1] - 0.320_627 * ictcp[2];

    let l_lin = eotf_st2084(l, 1.0);
    let m_lin = eotf_st2084(m, 1.0);
    let s_lin = eotf_st2084(s, 1.0);

    [
        (3.436_61 * l_lin - 2.506_45 * m_lin + 0.069_845_4 * s_lin).max(0.0),
        (-0.791_33 * l_lin + 1.983_6 * m_lin - 0.192_271 * s_lin).max(0.0),
        (-0.025_949_9 * l_lin - 0.098_913_7 * m_lin + 1.124_86 * s_lin).max(0.0),
    ]
}

#[inline]
fn rgb_to_jzazbz(rgb: [f32; 3]) -> [f32; 3] {
    // Input: linear Rec.2020, coefficients adjusted for linear Rec.2020
    let l = rgb[0] * 0.530_004 + rgb[1] * 0.355_704 + rgb[2] * 0.086_090;
    let m = rgb[0] * 0.289_388 + rgb[1] * 0.525_395 + rgb[2] * 0.157_481;
    let s = rgb[0] * 0.091_098 + rgb[1] * 0.147_588 + rgb[2] * 0.734_234;

    let l_pq = inverse_eotf_st2084(l, JZAZBZ_EXPONENT_SCALE_FACTOR);
    let m_pq = inverse_eotf_st2084(m, JZAZBZ_EXPONENT_SCALE_FACTOR);
    let s_pq = inverse_eotf_st2084(s, JZAZBZ_EXPONENT_SCALE_FACTOR);

    let iz = 0.5 * l_pq + 0.5 * m_pq;
    let jz = (0.44 * iz) / (1.0 - 0.56 * iz) - 1.629_549_9e-11;
    let az = 3.524_000 * l_pq - 4.066_708 * m_pq + 0.542_708 * s_pq;
    let bz = 0.199_076 * l_pq + 1.096_799 * m_pq - 1.295_875 * s_pq;
    [jz, az, bz]
}

#[inline]
fn jzazbz_to_rgb(jab: [f32; 3]) -> [f32; 3] {
    let jz = jab[0] + 1.629_549_9e-11;
    let iz = jz / (0.44 + 0.56 * jz);
    let a = jab[1];
    let b = jab[2];

    let l = iz + a * 1.386_050_5e-1 + b * 5.804_731_5e-2;
    let m = iz + a * -1.386_050_5e-1 + b * -5.804_731_5e-2;
    let s = iz + a * -9.601_924_4e-2 + b * -8.118_918_9e-1;

    let l_lin = eotf_st2084(l, JZAZBZ_EXPONENT_SCALE_FACTOR);
    let m_lin = eotf_st2084(m, JZAZBZ_EXPONENT_SCALE_FACTOR);
    let s_lin = eotf_st2084(s, JZAZBZ_EXPONENT_SCALE_FACTOR);

    [
        l_lin * 2.990_669 + m_lin * -2.049_742 + s_lin * 0.088_977,
        l_lin * -1.634_525 + m_lin * 3.145_627 + s_lin * -0.483_037,
        l_lin * -0.042_505 + m_lin * -0.377_983 + s_lin * 1.448_019,
    ]
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum UcsMode {
    ICtCp,
    Jzazbz,
}

#[inline]
fn rgb_to_ucs(rgb: [f32; 3], mode: UcsMode) -> [f32; 3] {
    match mode {
        UcsMode::ICtCp => rgb_to_ictcp(rgb),
        UcsMode::Jzazbz => rgb_to_jzazbz(rgb),
    }
}

#[inline]
fn ucs_to_rgb(ucs: [f32; 3], mode: UcsMode) -> [f32; 3] {
    match mode {
        UcsMode::ICtCp => ictcp_to_rgb(ucs),
        UcsMode::Jzazbz => jzazbz_to_rgb(ucs),
    }
}

#[derive(Debug, Clone)]
struct GTToneMappingCurveV2 {
    peak_intensity: f32,
    mid_point: f32,
    linear_section: f32,
    toe_strength: f32,
    k_a: f32,
    k_b: f32,
    k_c: f32,
}

impl GTToneMappingCurveV2 {
    fn new(
        monitor_intensity: f32,
        alpha: f32,
        gray_point: f32,
        linear_section: f32,
        toe_strength: f32,
    ) -> Self {
        let peak_intensity = monitor_intensity;
        let mid_point = gray_point;

        let k = (linear_section - 1.0) / (alpha - 1.0);
        let k_a = peak_intensity * linear_section + peak_intensity * k;
        let k_b = -peak_intensity * k * (linear_section / k).exp();
        let k_c = -1.0 / (k * peak_intensity);

        Self {
            peak_intensity,
            mid_point,
            linear_section,
            toe_strength,
            k_a,
            k_b,
            k_c,
        }
    }

    fn evaluate_curve(&self, x: f32) -> f32 {
        if x < 0.0 {
            return 0.0;
        }

        let weight_linear = smooth_step(x, 0.0, self.mid_point);
        let weight_toe = 1.0 - weight_linear;

        // Shoulder mapping for highlights
        let shoulder = self.k_a + self.k_b * (x * self.k_c).exp();

        if x < self.linear_section * self.peak_intensity {
            let toe_mapped = self.mid_point * (x / self.mid_point).powf(self.toe_strength);
            weight_toe * toe_mapped + weight_linear * x
        } else {
            shoulder
        }
    }
}

#[derive(Debug, Clone)]
pub struct GT7ToneMapping {
    pub ucs_mode: UcsMode,
    sdr_correction_factor: f32,
    framebuffer_luminance_target: f32,
    framebuffer_luminance_target_ucs: f32,
    curve: GTToneMappingCurveV2,
    blend_ratio: f32,
    fade_start: f32,
    fade_end: f32,
}

impl GT7ToneMapping {
    pub fn new(physical_target_luminance: Option<f32>, ucs_mode: UcsMode) -> Self {
        let (physical_target_luminance, sdr_correction_factor) =
            if let Some(l) = physical_target_luminance {
                (l, 1.0)
            } else {
                (
                    GRAN_TURISMO_SDR_PAPER_WHITE,
                    1.0 / physical_to_frame_buffer(GRAN_TURISMO_SDR_PAPER_WHITE),
                )
            };
        let framebuffer_luminance_target = physical_to_frame_buffer(physical_target_luminance);
        let curve =
            GTToneMappingCurveV2::new(framebuffer_luminance_target, 0.25, 0.538, 0.444, 1.280);

        // blend_ratio, fade_start, fade_end are already defaulted but keep parity with C++
        let blend_ratio = 0.6;
        let fade_start = 0.98;
        let fade_end = 1.16;

        let rgb = [framebuffer_luminance_target; 3];
        let ucs = rgb_to_ucs(rgb, ucs_mode);
        let framebuffer_luminance_target_ucs = ucs[0]; // luminance channel in UCS

        Self {
            ucs_mode,
            sdr_correction_factor,
            framebuffer_luminance_target,
            framebuffer_luminance_target_ucs,
            curve,
            blend_ratio,
            fade_start,
            fade_end,
        }
    }

    /// Apply tone mapping to linear Rec.2020 RGB frame-buffer values.
    /// Returns frame-buffer values ready for sRGB OETF (SDR) or PQ inverse-EOTF (HDR).
    pub fn apply(&self, rgb: [f32; 3]) -> [f32; 3] {
        // Convert to UCS to separate luminance and chroma.
        let ucs = rgb_to_ucs(rgb, self.ucs_mode);

        // Per-channel tone mapping (skewed color)
        let skewed_rgb = rgb.map(|x| self.curve.evaluate_curve(x));
        let skewed_ucs = rgb_to_ucs(skewed_rgb, self.ucs_mode);

        let chroma_scale = chroma_curve(
            ucs[0] / self.framebuffer_luminance_target_ucs,
            self.fade_start,
            self.fade_end,
        );

        let scaled_ucs = [skewed_ucs[0], ucs[1] * chroma_scale, ucs[2] * chroma_scale];

        let scaled_rgb = ucs_to_rgb(scaled_ucs, self.ucs_mode);

        let mut out = [0.0; 3];
        for i in 0..3 {
            let blended =
                (1.0 - self.blend_ratio) * skewed_rgb[i] + self.blend_ratio * scaled_rgb[i];
            out[i] = self.sdr_correction_factor * blended.min(self.framebuffer_luminance_target);
        }
        out
    }
}