edgefirst-tensor 0.25.2

// SPDX-FileCopyrightText: Copyright 2026 Au-Zone Technologies
// SPDX-License-Identifier: Apache-2.0

//! Colorimetry metadata for image/video tensors.
//!
//! Four orthogonal axes mirroring V4L2's `struct v4l2_format` colorimetry
//! fields and `libcamera::ColorSpace`, named to match the EdgeFirst
//! `CameraFrame.msg` schema so values round-trip through the ROS layer.
//! Each enum is `#[non_exhaustive]`; unknown/`_DEFAULT` values map to `None`.

use core::fmt;
use serde::{Deserialize, Serialize};

// V4L2 UAPI constants (stable kernel ABI) — mirrored from <linux/videodev2.h>.
const V4L2_COLORSPACE_SMPTE170M: u32 = 1;
const V4L2_COLORSPACE_REC709: u32 = 3;
const V4L2_COLORSPACE_470_SYSTEM_M: u32 = 5;
const V4L2_COLORSPACE_470_SYSTEM_BG: u32 = 6;
const V4L2_COLORSPACE_JPEG: u32 = 7;
const V4L2_COLORSPACE_SRGB: u32 = 8;
const V4L2_COLORSPACE_BT2020: u32 = 10;
const V4L2_XFER_FUNC_709: u32 = 1;
const V4L2_XFER_FUNC_SRGB: u32 = 2;
const V4L2_XFER_FUNC_NONE: u32 = 5;
const V4L2_XFER_FUNC_SMPTE2084: u32 = 7;
const V4L2_YCBCR_ENC_DEFAULT: u32 = 0;
const V4L2_YCBCR_ENC_601: u32 = 1;
const V4L2_YCBCR_ENC_709: u32 = 2;
const V4L2_YCBCR_ENC_BT2020: u32 = 6;
const V4L2_QUANTIZATION_DEFAULT: u32 = 0;
const V4L2_QUANTIZATION_FULL_RANGE: u32 = 1;
const V4L2_QUANTIZATION_LIM_RANGE: u32 = 2;

/// Color primaries (`color_space` in the EdgeFirst schema).
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ColorSpace {
    Bt709,
    Bt2020,
    Srgb,
    Smpte170m,
}

/// Transfer function (`color_transfer` in the EdgeFirst schema).
///
/// **Limitation:** this axis is stored, propagated, and round-tripped, but it is
/// **not applied** by any conversion backend. All YUV↔RGB paths operate only on
/// the matrix ([`ColorEncoding`]) and range ([`ColorRange`]); the transfer
/// function (gamma / TRC) is assumed to be the platform-native curve and is left
/// unchanged. HDR transfer curves ([`Self::Pq`] / [`Self::Hlg`]) are therefore
/// *not* tone-mapped — consumers needing linear-light or HDR handling must apply
/// the curve themselves.
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ColorTransfer {
    Bt709,
    Srgb,
    Pq,
    /// Hybrid Log-Gamma. Present for EdgeFirst-schema / libcamera parity; the
    /// V4L2 UAPI defines no `V4L2_XFER_FUNC_HLG`, so `from_v4l2` never yields
    /// this variant.
    Hlg,
    Linear,
}

/// YCbCr encoding matrix (`color_encoding` in the EdgeFirst schema).
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ColorEncoding {
    Bt601,
    Bt709,
    Bt2020,
}

/// Quantization range (`color_range` in the EdgeFirst schema).
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ColorRange {
    Full,
    Limited,
}

impl ColorSpace {
    /// Short string label matching the EdgeFirst schema.
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Bt709 => "bt709",
            Self::Bt2020 => "bt2020",
            Self::Srgb => "srgb",
            Self::Smpte170m => "smpte170m",
        }
    }

    /// Map a raw V4L2 `colorspace` field to a [`ColorSpace`].
    ///
    /// Returns `None` for `V4L2_COLORSPACE_DEFAULT` (0) and any
    /// unrecognised value.
    pub fn from_v4l2(v: u32) -> Option<Self> {
        match v {
            V4L2_COLORSPACE_SMPTE170M => Some(Self::Smpte170m),
            V4L2_COLORSPACE_REC709 => Some(Self::Bt709),
            // legacy NTSC (470M) / PAL-SECAM (470BG): close enough to SMPTE 170M primaries for conversion
            V4L2_COLORSPACE_470_SYSTEM_M | V4L2_COLORSPACE_470_SYSTEM_BG => Some(Self::Smpte170m),
            V4L2_COLORSPACE_JPEG | V4L2_COLORSPACE_SRGB => Some(Self::Srgb),
            V4L2_COLORSPACE_BT2020 => Some(Self::Bt2020),
            _ => None,
        }
    }
}

impl ColorTransfer {
    /// Short string label matching the EdgeFirst schema.
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Bt709 => "bt709",
            Self::Srgb => "srgb",
            Self::Pq => "pq",
            Self::Hlg => "hlg",
            Self::Linear => "linear",
        }
    }

    /// Map a raw V4L2 `xfer_func` field to a [`ColorTransfer`].
    ///
    /// Returns `None` for `V4L2_XFER_FUNC_DEFAULT` (0) and any
    /// unrecognised value.
    ///
    /// Note: OPRGB (3), SMPTE240M (4), and DCI_P3 (6) have no HAL equivalent
    /// and map to `None`. [`ColorTransfer::Hlg`] is never produced because the
    /// V4L2 UAPI defines no `V4L2_XFER_FUNC_HLG` value.
    pub fn from_v4l2(v: u32) -> Option<Self> {
        match v {
            V4L2_XFER_FUNC_709 => Some(Self::Bt709),
            V4L2_XFER_FUNC_SRGB => Some(Self::Srgb),
            V4L2_XFER_FUNC_NONE => Some(Self::Linear),
            V4L2_XFER_FUNC_SMPTE2084 => Some(Self::Pq),
            _ => None,
        }
    }
}

impl ColorEncoding {
    /// Short string label matching the EdgeFirst schema.
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Bt601 => "bt601",
            Self::Bt709 => "bt709",
            Self::Bt2020 => "bt2020",
        }
    }

    /// Map a raw V4L2 `ycbcr_enc` field to a [`ColorEncoding`].
    ///
    /// Returns `None` for `V4L2_YCBCR_ENC_DEFAULT` (0) and any
    /// unrecognised value.
    pub fn from_v4l2(v: u32) -> Option<Self> {
        match v {
            V4L2_YCBCR_ENC_601 => Some(Self::Bt601),
            V4L2_YCBCR_ENC_709 => Some(Self::Bt709),
            V4L2_YCBCR_ENC_BT2020 => Some(Self::Bt2020),
            _ => None,
        }
    }
}

impl ColorRange {
    /// Short string label matching the EdgeFirst schema.
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Full => "full",
            Self::Limited => "limited",
        }
    }

    /// Map a raw V4L2 `quantization` field to a [`ColorRange`].
    ///
    /// Returns `None` for `V4L2_QUANTIZATION_DEFAULT` (0) and any
    /// unrecognised value.
    pub fn from_v4l2(v: u32) -> Option<Self> {
        match v {
            V4L2_QUANTIZATION_FULL_RANGE => Some(Self::Full),
            V4L2_QUANTIZATION_LIM_RANGE => Some(Self::Limited),
            _ => None,
        }
    }
}

macro_rules! display_via_as_str {
    ($($t:ty),*) => {$(
        impl fmt::Display for $t {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                f.write_str(self.as_str())
            }
        }
    )*};
}
display_via_as_str!(ColorSpace, ColorTransfer, ColorEncoding, ColorRange);

/// YCbCr matrix luma weights `(kr, kb)`; the green weight is `kg = 1 - kr - kb`.
///
/// This is the single source of the BT.601/709/2020 coefficients. Every
/// YUV↔RGB path (the in-shader GL coefficients and the hand-rolled fixed-point
/// CPU encoders) derives its matrix from here so a coefficient change is made
/// in exactly one place.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct MatrixWeights {
    pub kr: f64,
    pub kb: f64,
}

impl MatrixWeights {
    /// Green luma weight, `1 - kr - kb`.
    pub fn kg(self) -> f64 {
        1.0 - self.kr - self.kb
    }
}

impl ColorEncoding {
    /// The matrix luma weights `(kr, kb)` for this encoding — the canonical
    /// BT.601/709/2020 coefficients shared by every YUV↔RGB conversion path.
    pub fn luma_weights(self) -> MatrixWeights {
        match self {
            Self::Bt601 => MatrixWeights {
                kr: 0.299,
                kb: 0.114,
            },
            Self::Bt709 => MatrixWeights {
                kr: 0.2126,
                kb: 0.0722,
            },
            Self::Bt2020 => MatrixWeights {
                kr: 0.2627,
                kb: 0.0593,
            },
        }
    }
}

/// Quantization swings (out of 255) for a YCbCr range: the luma black offset and
/// the luma/chroma excursions. Full range is `0 / 255 / 255`; limited (studio)
/// range is `16 / 219 / 224`.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct RangeScaling {
    /// Luma black level (0 for full range, 16 for limited).
    pub y_offset: f64,
    /// Luma excursion (255 for full range, 219 for limited).
    pub y_swing: f64,
    /// Chroma excursion (255 for full range, 224 for limited).
    pub c_swing: f64,
}

impl ColorRange {
    /// The luma/chroma quantization swings for this range — the canonical
    /// 0/255/255 (full) or 16/219/224 (limited) studio-swing constants shared
    /// by every YUV↔RGB conversion path.
    pub fn scaling(self) -> RangeScaling {
        match self {
            Self::Full => RangeScaling {
                y_offset: 0.0,
                y_swing: 255.0,
                c_swing: 255.0,
            },
            Self::Limited => RangeScaling {
                y_offset: 16.0,
                y_swing: 219.0,
                c_swing: 224.0,
            },
        }
    }
}

/// Full 4-axis colorimetry. Each axis is `Option`; `None` means "undefined"
/// and is never auto-filled — consumers (e.g. `convert()`) resolve missing
/// axes at use-time without mutating the value.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default, Serialize, Deserialize)]
pub struct Colorimetry {
    pub space: Option<ColorSpace>,
    pub transfer: Option<ColorTransfer>,
    pub encoding: Option<ColorEncoding>,
    pub range: Option<ColorRange>,
}

impl Colorimetry {
    /// JPEG/JFIF colorimetry: sRGB primaries, sRGB transfer, BT.601
    /// encoding, full range.
    pub fn jfif() -> Self {
        Self {
            space: Some(ColorSpace::Srgb),
            transfer: Some(ColorTransfer::Srgb),
            encoding: Some(ColorEncoding::Bt601),
            range: Some(ColorRange::Full),
        }
    }

    /// Build from the four raw V4L2 colorimetry integers.
    ///
    /// Explicit values map directly. For `ycbcr_enc`/`quantization`, the V4L2
    /// `DEFAULT` (0) sentinel does NOT mean "unknown" — it means "derive from
    /// the colorspace" (kernel `V4L2_MAP_YCBCR_ENC_DEFAULT` /
    /// `V4L2_MAP_QUANTIZATION_DEFAULT`). So a recognised colorspace resolves
    /// those axes here (e.g. `V4L2_COLORSPACE_JPEG` → BT.601 full-range) rather
    /// than leaving them `None` and falling through to the at-use height
    /// heuristic (which would wrongly pick BT.709/limited for an HD JPEG frame).
    /// A `DEFAULT`/unrecognised colorspace still yields `None` (deferred to the
    /// heuristic); unrecognised non-default values also map to `None`.
    pub fn from_v4l2(colorspace: u32, xfer: u32, ycbcr_enc: u32, quant: u32) -> Self {
        let encoding = ColorEncoding::from_v4l2(ycbcr_enc).or_else(|| {
            if ycbcr_enc == V4L2_YCBCR_ENC_DEFAULT {
                Self::default_encoding_for_colorspace(colorspace)
            } else {
                None
            }
        });
        let range = ColorRange::from_v4l2(quant).or_else(|| {
            if quant == V4L2_QUANTIZATION_DEFAULT {
                Self::default_range_for_colorspace(colorspace)
            } else {
                None
            }
        });
        Self {
            space: ColorSpace::from_v4l2(colorspace),
            transfer: ColorTransfer::from_v4l2(xfer),
            encoding,
            range,
        }
    }

    /// V4L2 `ycbcr_enc=DEFAULT` → encoding implied by the colorspace
    /// (`V4L2_MAP_YCBCR_ENC_DEFAULT`). `None` for default/unrecognised.
    fn default_encoding_for_colorspace(colorspace: u32) -> Option<ColorEncoding> {
        match colorspace {
            V4L2_COLORSPACE_REC709 => Some(ColorEncoding::Bt709),
            V4L2_COLORSPACE_BT2020 => Some(ColorEncoding::Bt2020),
            V4L2_COLORSPACE_SMPTE170M
            | V4L2_COLORSPACE_470_SYSTEM_M
            | V4L2_COLORSPACE_470_SYSTEM_BG
            | V4L2_COLORSPACE_JPEG
            | V4L2_COLORSPACE_SRGB => Some(ColorEncoding::Bt601),
            _ => None,
        }
    }

    /// V4L2 `quantization=DEFAULT` → range implied by the colorspace
    /// (`V4L2_MAP_QUANTIZATION_DEFAULT` for the YUV case: only JPEG is full,
    /// every other recognised colorspace is limited). `None` for
    /// default/unrecognised.
    fn default_range_for_colorspace(colorspace: u32) -> Option<ColorRange> {
        match colorspace {
            V4L2_COLORSPACE_JPEG => Some(ColorRange::Full),
            V4L2_COLORSPACE_SMPTE170M
            | V4L2_COLORSPACE_REC709
            | V4L2_COLORSPACE_BT2020
            | V4L2_COLORSPACE_470_SYSTEM_M
            | V4L2_COLORSPACE_470_SYSTEM_BG
            | V4L2_COLORSPACE_SRGB => Some(ColorRange::Limited),
            _ => None,
        }
    }

    /// Set color primaries (consuming builder).
    pub fn with_space(mut self, s: ColorSpace) -> Self {
        self.space = Some(s);
        self
    }

    /// Set transfer function (consuming builder).
    pub fn with_transfer(mut self, t: ColorTransfer) -> Self {
        self.transfer = Some(t);
        self
    }

    /// Set YCbCr encoding matrix (consuming builder).
    pub fn with_encoding(mut self, e: ColorEncoding) -> Self {
        self.encoding = Some(e);
        self
    }

    /// Set quantization range (consuming builder).
    pub fn with_range(mut self, r: ColorRange) -> Self {
        self.range = Some(r);
        self
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn from_v4l2_known_and_default() {
        assert_eq!(ColorEncoding::from_v4l2(1), Some(ColorEncoding::Bt601));
        assert_eq!(ColorEncoding::from_v4l2(2), Some(ColorEncoding::Bt709));
        assert_eq!(ColorEncoding::from_v4l2(6), Some(ColorEncoding::Bt2020));
        assert_eq!(ColorEncoding::from_v4l2(0), None); // DEFAULT
        assert_eq!(ColorEncoding::from_v4l2(3), None); // XV601 (unsurfaced)
        assert_eq!(ColorRange::from_v4l2(1), Some(ColorRange::Full));
        assert_eq!(ColorRange::from_v4l2(2), Some(ColorRange::Limited));
        assert_eq!(ColorSpace::from_v4l2(7), Some(ColorSpace::Srgb)); // JPEG→sRGB
        assert_eq!(ColorTransfer::from_v4l2(5), Some(ColorTransfer::Linear)); // NONE→linear

        // ColorSpace: additional arms
        assert_eq!(ColorSpace::from_v4l2(1), Some(ColorSpace::Smpte170m)); // SMPTE170M
        assert_eq!(ColorSpace::from_v4l2(3), Some(ColorSpace::Bt709)); // REC709
        assert_eq!(ColorSpace::from_v4l2(10), Some(ColorSpace::Bt2020)); // BT2020
        assert_eq!(ColorSpace::from_v4l2(5), Some(ColorSpace::Smpte170m)); // 470_SYSTEM_M
        assert_eq!(ColorSpace::from_v4l2(6), Some(ColorSpace::Smpte170m)); // 470_SYSTEM_BG

        // ColorTransfer: additional arms and unmapped values
        assert_eq!(ColorTransfer::from_v4l2(1), Some(ColorTransfer::Bt709)); // XFER_FUNC_709
        assert_eq!(ColorTransfer::from_v4l2(2), Some(ColorTransfer::Srgb)); // XFER_FUNC_SRGB
        assert_eq!(ColorTransfer::from_v4l2(7), Some(ColorTransfer::Pq)); // XFER_FUNC_SMPTE2084
        assert_eq!(ColorTransfer::from_v4l2(3), None); // OPRGB — no HAL equivalent
        assert_eq!(ColorTransfer::from_v4l2(6), None); // DCI_P3 — no HAL equivalent

        // Hlg is never produced by from_v4l2 — no V4L2_XFER_FUNC_HLG exists in the kernel UAPI
        for v in 0u32..=10 {
            assert_ne!(ColorTransfer::from_v4l2(v), Some(ColorTransfer::Hlg));
        }
    }

    #[test]
    fn jfif_is_bt601_full_srgb() {
        let c = Colorimetry::jfif();
        assert_eq!(c.space, Some(ColorSpace::Srgb));
        assert_eq!(c.transfer, Some(ColorTransfer::Srgb));
        assert_eq!(c.encoding, Some(ColorEncoding::Bt601));
        assert_eq!(c.range, Some(ColorRange::Full));
    }

    #[test]
    fn from_v4l2_struct_maps_all_axes_and_unknown_to_none() {
        let c = Colorimetry::from_v4l2(3, 1, 2, 1); // REC709, XFER709, ENC709, FULL
        assert_eq!(c.space, Some(ColorSpace::Bt709));
        assert_eq!(c.transfer, Some(ColorTransfer::Bt709));
        assert_eq!(c.encoding, Some(ColorEncoding::Bt709));
        assert_eq!(c.range, Some(ColorRange::Full));
        let d = Colorimetry::from_v4l2(0, 0, 0, 0); // all DEFAULT
        assert_eq!(d, Colorimetry::default()); // all None
    }

    #[test]
    fn from_v4l2_default_enc_quant_derive_from_colorspace() {
        // COLORSPACE_JPEG (7) with DEFAULT ycbcr_enc/quant must resolve to
        // BT.601 full-range per V4L2_MAP_*_DEFAULT — NOT be left None (which
        // would let the height heuristic wrongly pick BT.709/limited for HD).
        let jpeg = Colorimetry::from_v4l2(7, 0, 0, 0);
        assert_eq!(jpeg.encoding, Some(ColorEncoding::Bt601));
        assert_eq!(jpeg.range, Some(ColorRange::Full));

        // REC709 colorspace, DEFAULT enc/quant → BT.709 limited.
        let rec709 = Colorimetry::from_v4l2(3, 0, 0, 0);
        assert_eq!(rec709.encoding, Some(ColorEncoding::Bt709));
        assert_eq!(rec709.range, Some(ColorRange::Limited));

        // Explicit ycbcr_enc/quant still win over the colorspace default.
        let explicit = Colorimetry::from_v4l2(7, 0, 2, 2); // JPEG but enc=709, quant=limited
        assert_eq!(explicit.encoding, Some(ColorEncoding::Bt709));
        assert_eq!(explicit.range, Some(ColorRange::Limited));

        // Unrecognised non-default values stay None (not derived).
        let unknown_enc = Colorimetry::from_v4l2(7, 0, 99, 0);
        assert_eq!(unknown_enc.encoding, None);
        assert_eq!(unknown_enc.range, Some(ColorRange::Full)); // quant still defaulted from JPEG
    }

    #[test]
    fn luma_weights_are_the_canonical_bt_constants() {
        assert_eq!(
            ColorEncoding::Bt601.luma_weights(),
            MatrixWeights {
                kr: 0.299,
                kb: 0.114
            }
        );
        assert_eq!(
            ColorEncoding::Bt709.luma_weights(),
            MatrixWeights {
                kr: 0.2126,
                kb: 0.0722
            }
        );
        assert_eq!(
            ColorEncoding::Bt2020.luma_weights(),
            MatrixWeights {
                kr: 0.2627,
                kb: 0.0593
            }
        );
        // kg = 1 - kr - kb (BT.709 green weight ≈ 0.7152).
        assert!((ColorEncoding::Bt709.luma_weights().kg() - 0.7152).abs() < 1e-9);
    }

    #[test]
    fn range_scaling_is_full_or_studio_swing() {
        let full = ColorRange::Full.scaling();
        assert_eq!(
            full,
            RangeScaling {
                y_offset: 0.0,
                y_swing: 255.0,
                c_swing: 255.0
            }
        );
        let limited = ColorRange::Limited.scaling();
        assert_eq!(
            limited,
            RangeScaling {
                y_offset: 16.0,
                y_swing: 219.0,
                c_swing: 224.0
            }
        );
    }

    #[test]
    fn as_str_matches_schema() {
        assert_eq!(ColorEncoding::Bt601.as_str(), "bt601");
        assert_eq!(ColorRange::Full.as_str(), "full");
        assert_eq!(ColorSpace::Smpte170m.as_str(), "smpte170m");
        assert_eq!(ColorTransfer::Pq.as_str(), "pq");
        assert_eq!(ColorTransfer::Hlg.as_str(), "hlg");
    }
}