diff --git a/avtensor.pyi b/avtensor.pyi index 4d7e321..ff157e9 100644 --- a/avtensor.pyi +++ b/avtensor.pyi @@ -1,6 +1,6 @@ import sys from enum import Enum -from typing import TypedDict +from typing import Literal, TypedDict if sys.version_info < (3, 11): from typing_extensions import NotRequired @@ -32,6 +32,12 @@ class VideoStreamRequest: # "uint8" (default) or "float32". float32 decodes to planar float in # [0, 1] (NCHW-contiguous), preserving the depth of 10/12-bit sources. dtype: str | None + # HDR handling for PQ/HLG or wide-gamut sources: "tonemap" (default) + # tone maps to an SDR BT.709 preview; "raw" preserves the source's code + # values (tagged matrix/range only — transfer function untouched). Use + # "raw" when you need the actual HDR signal, e.g. training on PQ + # masters or colorimetric measurement. + hdr_mode: Literal["tonemap", "raw"] | None def __init__( self, @@ -45,6 +51,7 @@ class VideoStreamRequest: dimension_order: str | None = None, device: str | None = None, dtype: str | None = None, + hdr_mode: Literal["tonemap", "raw"] | None = None, ): ... class LoudnessNormalization: diff --git a/src/decoder/mod.rs b/src/decoder/mod.rs index 07a4221..2f1f92f 100644 --- a/src/decoder/mod.rs +++ b/src/decoder/mod.rs @@ -150,6 +150,9 @@ pub struct VideoStreamRequest { /// Element type of the decoded video tensor. `Float32` decodes via /// 16-bit RGB, preserving the full depth of >8-bit sources. pub dtype: OutputDtype, + /// HDR handling: `Tonemap` (default, SDR BT.709 preview) or `Raw` + /// (preserve PQ/HLG code values; matrix/range conversion only). + pub hdr_mode: HdrMode, } #[derive(Default, Debug, Clone)] @@ -2191,6 +2194,35 @@ impl SourceColorInfo { } } +/// How HDR/wide-gamut sources are handled during decode. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum HdrMode { + /// Tone map to SDR BT.709 (the historical default): linearize, hable + /// tone map, convert primaries/transfer/matrix to BT.709. + #[default] + Tonemap, + /// Preserve the source's raw code values: YUV→RGB uses the stream's + /// tagged matrix and range only; the transfer function (e.g. PQ/HLG) + /// is NOT linearized, tone mapped, or converted. Required when the + /// consumer needs the actual HDR signal (training on PQ masters, + /// colorimetric measurement) rather than an SDR preview. + Raw, +} + +impl TryFrom> for HdrMode { + type Error = anyhow::Error; + + fn try_from(value: Option<&str>) -> Result { + match value { + None | Some("tonemap") => Ok(Self::Tonemap), + Some("raw") => Ok(Self::Raw), + Some(other) => Err(anyhow::anyhow!( + "hdr_mode must be \"tonemap\" or \"raw\", got {other:?}" + )), + } + } +} + #[derive(Debug)] pub struct VideoFilterConfig { /// Desired frame rate for the video. @@ -2205,6 +2237,8 @@ pub struct VideoFilterConfig { source_color: SourceColorInfo, /// CUDA ordinal for GPU-resident output (frames stay on the GPU). device: Option, + /// How HDR sources are converted (tone-mapped SDR vs raw code values). + hdr_mode: HdrMode, } impl Display for VideoFilterConfig { @@ -2227,7 +2261,10 @@ impl Display for VideoFilterConfig { // For HDR/wide-gamut content, insert tone mapping and gamut conversion // to produce correct sRGB output instead of a naive YUV→RGB conversion. - if self.source_color.is_hdr() { + // With `HdrMode::Raw` the block is skipped entirely: the plain + // format conversion below honors the tagged matrix/range but leaves + // the transfer function (PQ/HLG code values) untouched. + if self.source_color.is_hdr() && self.hdr_mode == HdrMode::Tonemap { log::debug!( "HDR source detected (trc={:?}, primaries={:?}), inserting tone mapping pipeline", self.source_color.color_trc, @@ -2261,6 +2298,7 @@ impl Default for VideoFilterConfig { height: Default::default(), source_color: Default::default(), device: Default::default(), + hdr_mode: Default::default(), } } } @@ -2278,6 +2316,7 @@ impl TryFrom<&VideoStreamRequest> for VideoFilterConfig { OutputDtype::Uint8 => "rgb24".to_string(), OutputDtype::Float32 => "gbrpf32le".to_string(), }, + hdr_mode: req.hdr_mode, ..Default::default() }) } @@ -3526,6 +3565,22 @@ mod tests { "zscale=t=linear:npl=100,format=gbrpf32le,tonemap=hable:desat=0,zscale=p=bt709:t=bt709:m=bt709:range=tv,format=pix_fmts=rgb24"; "PQ BT.2020 source triggers HDR tone mapping pipeline" )] + #[test_case( + VideoFilterConfig { + pixel_format: "rgb24".to_string(), + device: None, + source_color: SourceColorInfo { + color_trc: Some("smpte2084".to_string()), + color_primaries: Some("bt2020".to_string()), + colorspace: Some("bt2020nc".to_string()), + color_range: Some("tv".to_string()), + }, + hdr_mode: HdrMode::Raw, + ..Default::default() + }, + "format=pix_fmts=rgb24"; + "PQ BT.2020 source with raw hdr_mode skips tone mapping and keeps code values" + )] #[test_case( VideoFilterConfig { frame_rate: Some(24.0), @@ -3538,6 +3593,7 @@ mod tests { color_primaries: Some("bt2020".to_string()), ..Default::default() }, + hdr_mode: HdrMode::Tonemap, }, "fps=24,scale=width=1920:height=1080,zscale=t=linear:npl=100,format=gbrpf32le,tonemap=hable:desat=0,zscale=p=bt709:t=bt709:m=bt709:range=tv,format=pix_fmts=rgb24"; "HDR with fps and scale options" diff --git a/src/ffi/mod.rs b/src/ffi/mod.rs index acb05f8..b230249 100644 --- a/src/ffi/mod.rs +++ b/src/ffi/mod.rs @@ -89,6 +89,15 @@ pub struct VideoStreamRequest { /// of 10/12-bit sources instead of quantizing them to 8 bits. #[pyo3(get, set)] dtype: Option, + /// HDR handling for PQ/HLG or wide-gamut sources: "tonemap" (the + /// default) tone maps to an SDR BT.709 preview; "raw" preserves the + /// source's code values — YUV→RGB uses the tagged matrix/range only + /// and the transfer function is left untouched. Use "raw" whenever the + /// consumer needs the actual HDR signal (training on PQ masters, + /// colorimetric measurement); the tone-mapped default is display- + /// oriented and substantially alters both luminance and chroma. + #[pyo3(get, set)] + hdr_mode: Option, } impl VideoStreamRequest { @@ -108,7 +117,7 @@ impl VideoStreamRequest { #[pymethods] impl VideoStreamRequest { #[new] - #[pyo3(signature = (*, index=None, width=None, height=None, fps=None, number_of_threads=None, hardware_acceleration=None, dimension_order=None, device=None, dtype=None))] + #[pyo3(signature = (*, index=None, width=None, height=None, fps=None, number_of_threads=None, hardware_acceleration=None, dimension_order=None, device=None, dtype=None, hdr_mode=None))] #[allow(clippy::too_many_arguments)] pub fn py_new( index: Option, @@ -120,6 +129,7 @@ impl VideoStreamRequest { dimension_order: Option, device: Option, dtype: Option, + hdr_mode: Option, ) -> Self { VideoStreamRequest { index, @@ -131,6 +141,7 @@ impl VideoStreamRequest { dimension_order, device, dtype, + hdr_mode, } } } @@ -171,6 +182,7 @@ impl VideoStreamRequest { hardware_acceleration: self.hardware_acceleration, device: self.device_ordinal()?, dtype: self.dtype_parsed()?, + hdr_mode: decoder::HdrMode::try_from(self.hdr_mode.as_deref())?, }) } }