Add hdr_mode option: raw code-value decode for HDR sources

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:

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<Option<&str>> for HdrMode {
+    type Error = anyhow::Error;
+
+    fn try_from(value: Option<&str>) -> Result<Self, Self::Error> {
+        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<i32>,
+    /// 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"

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<String>,
+    /// 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<String>,
 }
 
 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<usize>,
@@ -120,6 +129,7 @@ impl VideoStreamRequest {
         dimension_order: Option<String>,
         device: Option<String>,
         dtype: Option<String>,
+        hdr_mode: Option<String>,
     ) -> 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())?,
         })
     }
 }