diff --git a/README.md b/README.md index 9ed6893..c4b2e23 100644 --- a/README.md +++ b/README.md @@ -185,30 +185,66 @@ CPU is the consumer's job; `examples/common` shows one way. ### Color Conversion (RGB → YUV) -PixelForge includes a GPU compute shader for converting RGB input to YUV output, supporting multiple color spaces: - -| Color Space | Description | -|-------------|-------------| -| `Bt709` | Standard SDR (BT.709 coefficients) | -| `Bt2020` | HDR passthrough (BT.2020 coefficients, PQ-encoded input) | -| `SrgbToBt2020Pq` | SDR-in-HDR (sRGB → linear → BT.2020 gamut → PQ OETF) | -| `Bt709LinearToBt2020Pq` | scRGB HDR (linear BT.709 → BT.2020 gamut → PQ OETF). `sdr_reference_white_nits` sets the interpretation of 1.0; per the scRGB spec (IEC 61966-2-2), 80 nits. | +PixelForge includes a GPU compute shader for converting RGB input to YUV +output. The source color describes the input. The target color describes the stream. +Both are a [`ColorSpec`]. + +| `ColorSpec` | | Can be a target | +|-------------|-|-----------------| +| `Srgb` | Ordinary SDR content | yes | +| `Bt709Linear` | scRGB, from an `EXTENDED_SRGB_LINEAR_EXT` swapchain | no | +| `Bt2020Linear` | Linear light, wide gamut | no | +| `Bt2020Pq` | HDR10, from an `HDR10_ST2084_EXT` swapchain | yes | + +The linear spaces cannot be a target, because a video file has no way to +record that it holds linear light. Any source can be converted to +`Bt2020Pq`. Only `Srgb` can be converted to `Srgb`; the others would need +tone mapping or a gamma curve applied, which the shader does not do. +[`ColorConverter::new`] rejects the combinations it cannot do. + +Encoding to HDR needs to know how bright the source's white is, since HDR +carries real brightness values and SDR does not. Each space has a sensible +default, see [`ColorSpec::reference_white_nits`], overridable with +[`ColorConverterConfig::with_reference_white_nits`]. Supported input formats: BGRx, RGBx, BGRA, RGBA, ABGR2101010 (10-bit packed), RGBA16F (FP16). Supported output formats: NV12 (8-bit), I420 (8-bit), YUV444 (8-bit), P010 (10-bit), YUV444P10 (10-bit). +Pass [`ColorConverter::color_description`] to the encoder rather than +describing the colours a second time by hand. The two have to agree: if the +converter writes full-range pixels and the stream says limited, players +stretch the range again and the picture comes out wrong. + ```rust -use pixelforge::{ColorConverter, ColorConverterConfig, ColorSpace, InputFormat, OutputFormat, VideoContextBuilder}; +use pixelforge::{ + Codec, ColorConverter, ColorConverterConfig, ColorRange, ColorSpec, EncodeConfig, + Encoder, InputFormat, OutputFormat, VideoContextBuilder, +}; let context = VideoContextBuilder::new() .app_name("Color Converter") + .require_encode(Codec::H265) .build()?; -let mut config = ColorConverterConfig::new(1920, 1080, InputFormat::BGRx, OutputFormat::NV12); -config.color_space = ColorSpace::SrgbToBt2020Pq; +// SDR desktop content, encoded as HDR10 for an HDR streaming session. +let config = ColorConverterConfig::new( + 1920, + 1080, + InputFormat::BGRx, + OutputFormat::P010, + ColorSpec::Srgb, + ColorSpec::Bt2020Pq, + ColorRange::Full, +); let mut converter = ColorConverter::new(context.clone(), config)?; -// converter.convert(input_image, output_buffer)?; + +// The encoder declares exactly what the shader wrote. +let encode_config = EncodeConfig::h265(1920, 1080) + .with_color_description(converter.color_description()); + +let mut encoder = Encoder::new(context, encode_config)?; +// converter.convert(input_image, layout, encoder.input_image())?; ``` ## Benchmarking diff --git a/shader/color_convert.comp b/shader/color_convert.comp index f726ef8..c9f2c92 100644 --- a/shader/color_convert.comp +++ b/shader/color_convert.comp @@ -7,11 +7,19 @@ layout(push_constant) uniform PushConstants { uint height; uint input_format; // 0=BGRx, 1=RGBx, 2=BGRA, 3=RGBA, 4=ABGR2101010, 5=RGBA16F uint output_format; // 0=NV12, 1=I420, 2=YUV444, 3=P010, 4=YUV444P10 - uint color_space; // 0=BT.709, 1=BT.2020, 2=sRGB→BT.2020+PQ, 3=scRGB linear→BT.2020+PQ - uint full_range; // 0=limited/studio range, 1=full range - float sdr_white_nits; // SDR reference white (nits), used for sRGB and scRGB → BT.2020+PQ + uint source; // what the input is, as a ColorSpec discriminant + uint target; // what to produce, as a ColorSpec discriminant + uint range; // 0=limited/studio range, 1=full range + float reference_white_nits; // luminance of source value 1.0; unused when the source is PQ } params; +// ColorSpec discriminants. One set: `source` and `target` are the same type, +// and which end a value describes comes from which push constant it is in. +#define SPEC_SRGB 0u +#define SPEC_BT709_LINEAR 1u +#define SPEC_BT2020_LINEAR 2u +#define SPEC_BT2020_PQ 3u + // Source image sampled directly — eliminates the image-to-buffer copy. layout(binding = 0) uniform sampler2D inputImage; @@ -89,25 +97,26 @@ vec3 bt709_to_bt2020(vec3 rgb709) { // - PQ surfaces → PQ-encoded values (even in FP16) // - sRGB surfaces → sRGB gamma-encoded values (even in FP16) // -// color_space 0 (BT.709) and 1 (BT.2020): passthrough — data is already encoded. -// color_space 2 (sRGB→BT.2020+PQ): decode sRGB, convert gamut, apply PQ. -// color_space 3 (scRGB→BT.2020+PQ): convert gamut, apply PQ. No EOTF — input is linear. +// which is what `source` describes. vec3 read_rgb(ivec2 coord) { vec4 rgba = texelFetch(inputImage, coord, 0); - if (params.color_space == 2u) { - // sRGB→BT.2020+PQ: decode sRGB gamma → linear BT.709 → BT.2020 gamut → PQ. - vec3 linear_709 = srgb_to_linear(rgba.rgb); - vec3 linear_2020 = bt709_to_bt2020(linear_709); - // SDR reference white (configurable, default 203 nits per ITU-R BT.2408) - // normalized to PQ's 10000 nit scale. - return linear_to_pq(linear_2020 * (params.sdr_white_nits / 10000.0)); - } - if (params.color_space == 3u) { - // scRGB→BT.2020+PQ: BT.709 linear → BT.2020 gamut → PQ (no EOTF, input is already linear). - vec3 linear_2020 = bt709_to_bt2020(rgba.rgb); - return linear_to_pq(linear_2020 * (params.sdr_white_nits / 10000.0)); + + if (params.target == SPEC_BT2020_PQ) { + // Already PQ in the right gamut: nothing to do. + if (params.source == SPEC_BT2020_PQ) return rgba.rgb; + + // Get to linear light, then into BT.2020, then encode PQ. The + // reference white says what luminance source 1.0 stands for, which is + // what makes the PQ encode absolute. + vec3 linear = rgba.rgb; + if (params.source == SPEC_SRGB) linear = srgb_to_linear(linear); + if (params.source != SPEC_BT2020_LINEAR) linear = bt709_to_bt2020(linear); + return linear_to_pq(linear * (params.reference_white_nits / 10000.0)); } - // BT.709 or BT.2020 passthrough: values are already properly encoded. + + // An SDR target: the only source that reaches it is already SDR-encoded in + // the right gamut, so this is a passthrough. Every other source is + // rejected on the CPU side before the shader ever runs. return rgba.rgb; } @@ -115,8 +124,8 @@ vec3 read_rgb(ivec2 coord) { // Returns Y in [0, 1], U and V in [0, 1] centered at 0.5. vec3 rgb_to_yuv(vec3 rgb) { float yr, yg, yb, ur, ug, ub, vr, vg, vb; - if (params.color_space >= 1u) { - // BT.2020 (also used for sRGB→BT.2020+PQ, since read_rgb already converted). + if (params.target == SPEC_BT2020_PQ) { + // BT.2020 NCL. read_rgb has already put the samples in that gamut. yr = BT2020_Y_R; yg = BT2020_Y_G; yb = BT2020_Y_B; ur = BT2020_U_R; ug = BT2020_U_G; ub = BT2020_U_B; vr = BT2020_V_R; vg = BT2020_V_G; vb = BT2020_V_B; @@ -135,12 +144,12 @@ vec3 rgb_to_yuv(vec3 rgb) { // --- 8-bit quantization --- uint q8_y(float y) { - if (params.full_range == 0u) return uint(round(clamp(y * 219.0 + 16.0, 0.0, 255.0))); + if (params.range == 0u) return uint(round(clamp(y * 219.0 + 16.0, 0.0, 255.0))); return uint(round(clamp(y * 255.0, 0.0, 255.0))); } uint q8_c(float c) { - if (params.full_range == 0u) return uint(round(clamp((c - 0.5) * 224.0 + 128.0, 0.0, 255.0))); + if (params.range == 0u) return uint(round(clamp((c - 0.5) * 224.0 + 128.0, 0.0, 255.0))); return uint(round(clamp(c * 255.0, 0.0, 255.0))); } @@ -148,14 +157,14 @@ uint q8_c(float c) { uint q10_y(float y) { uint val; - if (params.full_range == 0u) val = uint(round(clamp(y * 876.0 + 64.0, 0.0, 1023.0))); + if (params.range == 0u) val = uint(round(clamp(y * 876.0 + 64.0, 0.0, 1023.0))); else val = uint(round(clamp(y * 1023.0, 0.0, 1023.0))); return (val << 6u) & 0xFFC0u; } uint q10_c(float c) { uint val; - if (params.full_range == 0u) val = uint(round(clamp((c - 0.5) * 896.0 + 512.0, 0.0, 1023.0))); + if (params.range == 0u) val = uint(round(clamp((c - 0.5) * 896.0 + 512.0, 0.0, 1023.0))); else val = uint(round(clamp(c * 1023.0, 0.0, 1023.0))); return (val << 6u) & 0xFFC0u; } diff --git a/shader/color_convert.spv b/shader/color_convert.spv index f0ad72c..6abab7f 100644 Binary files a/shader/color_convert.spv and b/shader/color_convert.spv differ diff --git a/src/converter/mod.rs b/src/converter/mod.rs index 770a894..dbabfea 100644 --- a/src/converter/mod.rs +++ b/src/converter/mod.rs @@ -8,27 +8,98 @@ mod pipeline; +use crate::encoder::ColorDescription; use crate::error::{PixelForgeError, Result}; use crate::vulkan::VideoContext; use ash::vk; use tracing::debug; -/// Color space for RGB→YUV conversion matrix selection. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] -pub enum ColorSpace { - /// BT.709 (standard SDR). Used for all HD/UHD SDR content. - #[default] - Bt709, - /// BT.2020 (HDR / wide color gamut). - Bt2020, - /// sRGB input converted to BT.2020 + PQ output. - /// Used for SDR content during HDR streaming sessions. - /// Applies: inverse sRGB EOTF → BT.709→BT.2020 gamut mapping → PQ OETF. - SrgbToBt2020Pq, - /// scRGB input (linear BT.709, FP16) converted to BT.2020 + PQ output. - /// Used for HDR games submitting a `VK_COLOR_SPACE_EXTENDED_SRGB_LINEAR_EXT` swapchain. - /// Applies: BT.709→BT.2020 gamut mapping → PQ OETF. No EOTF decode — input is already linear. - Bt709LinearToBt2020Pq, +/// A colour space, used for both ends of a conversion: +/// [`ColorConverterConfig::source`] for the input and +/// [`ColorConverterConfig::target`] for the encoded output. +/// +/// The two linear spaces can only be an input. Video files have no way to say +/// "linear", so nothing could be decoded correctly; see +/// [`ColorSpec::is_encodable`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +#[repr(u32)] +pub enum ColorSpec { + /// Ordinary SDR content, the usual desktop and game output. + /// + /// sRGB and BT.709 cover the same colours and differ only in the curve + /// applied to them, which is why the linear one below is named for BT.709 + /// rather than for sRGB. + Srgb = 0, + /// scRGB: the same colours as `Srgb`, in linear light. + /// + /// Values above 1.0 are allowed and mean brighter than SDR white, which is + /// how this carries HDR. What a `VK_COLOR_SPACE_EXTENDED_SRGB_LINEAR_EXT` + /// swapchain gives you. + Bt709Linear = 1, + /// Like `Bt709Linear`, but with the wider HDR colours. + Bt2020Linear = 2, + /// HDR10: wide gamut, PQ encoded. What a + /// `VK_COLOR_SPACE_HDR10_ST2084_EXT` swapchain gives you. + Bt2020Pq = 3, +} + +impl ColorSpec { + /// How bright this space's white is, in nits. + /// + /// Only used when encoding to HDR, which needs real brightness values + /// rather than the relative ones SDR content carries. Getting it wrong + /// makes the picture too bright or too dim, so the spaces differ here: + /// scRGB defines its white as 80 nits, the others use 203. + /// + /// `None` for `Bt2020Pq`, which already carries real brightness values. + /// Override it with + /// [`ColorConverterConfig::with_reference_white_nits`]. + pub fn reference_white_nits(&self) -> Option { + match self { + Self::Srgb | Self::Bt2020Linear => Some(203.0), + Self::Bt709Linear => Some(80.0), + Self::Bt2020Pq => None, + } + } + + /// Whether video can be encoded in this space. + /// + /// False for the two linear spaces: a video file has no way to record that + /// it holds linear light, so a player would get it wrong. + pub fn is_encodable(&self) -> bool { + matches!(self, Self::Srgb | Self::Bt2020Pq) + } + + /// What to tell a player about video in this space. + /// + /// `None` when [`Self::is_encodable`] is false. The range is set + /// separately, by [`ColorConverterConfig::color_description`]. + pub fn color_description(&self) -> Option { + match self { + // sRGB and BT.709 differ slightly in their curves, but video + // always labels this case BT.709 and players treat the two the + // same, so labelling it anything else would be surprising. + Self::Srgb => Some(ColorDescription::bt709()), + Self::Bt2020Pq => Some(ColorDescription::bt2020_pq()), + Self::Bt709Linear | Self::Bt2020Linear => None, + } + } +} + +/// Luma and chroma quantization range. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ColorRange { + /// Studio range: luma 16..235, chroma 16..240 at 8-bit. + Limited, + /// Full range: 0..255 at 8-bit. + Full, +} + +impl ColorRange { + /// Whether this is [`ColorRange::Full`]. + pub fn is_full(&self) -> bool { + matches!(self, Self::Full) + } } /// Supported input pixel formats for color conversion. @@ -49,14 +120,11 @@ pub enum InputFormat { /// RGBA16F (64-bit, 16-bit float per channel). /// Maps to DRM_FORMAT_ABGR16161616F / VK_FORMAT_R16G16B16A16_SFLOAT. /// - /// The converter treats FP16 data the same as other formats: passthrough - /// for `ColorSpace::Bt709` / `Bt2020`, sRGB→BT.2020+PQ conversion for - /// `ColorSpace::SrgbToBt2020Pq`, and scRGB-linear→BT.2020+PQ for - /// `ColorSpace::Bt709LinearToBt2020Pq` (the natural pairing for FP16 input - /// from an `EXTENDED_SRGB_LINEAR_EXT` swapchain). No linear→PQ transfer - /// function is applied automatically; if the source is PQ-encoded - /// (e.g. via the gamescope WSI layer), use `ColorSpace::Bt2020` for - /// passthrough. + /// The converter treats FP16 data the same as other formats: the + /// [`ColorSpec`] says what the samples mean, not the pixel format. + /// [`ColorSpec::Bt709Linear`] is the natural pairing for FP16 from an + /// `EXTENDED_SRGB_LINEAR_EXT` swapchain; FP16 that is already PQ-encoded + /// (via the gamescope WSI layer, say) is [`ColorSpec::Bt2020Pq`]. RGBA16F, } @@ -145,6 +213,9 @@ impl OutputFormat { } /// Configuration for the color converter. +/// +/// Colour is two decisions: [`source`](Self::source) for what the input pixels +/// are, [`target`](Self::target) for what the encoded stream should be. #[derive(Clone, Debug)] #[non_exhaustive] pub struct ColorConverterConfig { @@ -156,38 +227,130 @@ pub struct ColorConverterConfig { pub input_format: InputFormat, /// Output YUV format. pub output_format: OutputFormat, - /// Color space for the RGB→YUV matrix. - /// Use Bt709 for SDR, Bt2020 for HDR. - pub color_space: ColorSpace, - /// Full range (0-255 luma) or limited/studio range (16-235 luma). - /// Must match the `full_range` flag in `ColorDescription` for correct playback. - pub full_range: bool, - /// SDR reference white level in nits for the sRGB→BT.2020+PQ conversion. + /// What the input pixels already are. + pub source: ColorSpec, + /// What the encoded stream should be. Also picks the YUV matrix. /// - /// Per ITU-R BT.2408 the standard value is 203 nits. - /// Only used when `color_space` is `SrgbToBt2020Pq` or `Bt709LinearToBt2020Pq` - /// (for the latter, the scRGB spec defines 1.0 == 80 cd/m²). - pub sdr_reference_white_nits: f32, + /// Must be a space video can be encoded in, see + /// [`ColorSpec::is_encodable`]. + pub target: ColorSpec, + /// How bright the source's white is, in nits, overriding the source + /// space's default. Only used when encoding to HDR. + pub reference_white_nits: Option, + /// Whether the shader writes full or limited range. + /// + /// The stream has to say the same thing, which is what + /// [`color_description`](Self::color_description) is for. + pub range: ColorRange, } impl ColorConverterConfig { - /// Create a new configuration with BT.709 color space and limited range. + /// Create a new configuration. pub fn new( width: u32, height: u32, input_format: InputFormat, output_format: OutputFormat, + source_color: ColorSpec, + target_color: ColorSpec, + color_range: ColorRange, ) -> Self { Self { width, height, input_format, output_format, - color_space: ColorSpace::Bt709, - full_range: false, - sdr_reference_white_nits: 203.0, + source: source_color, + target: target_color, + range: color_range, + reference_white_nits: None, + } + } + + /// Override how bright the source's white is, in nits. + /// + /// Rarely needed; each space already has a sensible default, see + /// [`ColorSpec::reference_white_nits`]. + pub fn with_reference_white_nits(mut self, nits: f32) -> Self { + self.reference_white_nits = Some(nits); + self + } + + /// The value the shader gets: the override, or the source's default. + fn effective_reference_white_nits(&self) -> f32 { + self.reference_white_nits + .or_else(|| self.source.reference_white_nits()) + .unwrap_or(0.0) + } + + /// What to tell the encoder about the colours this conversion produces, + /// for + /// [`EncodeConfig::with_color_description`](crate::EncodeConfig::with_color_description). + /// + /// Taken from [`target`](Self::target) and [`range`](Self::range), so the + /// stream cannot end up describing something the shader did not write. + /// + /// `None` when the target is one [`ColorConverter::new`] would reject + /// anyway. [`ColorConverter::color_description`] gives the same answer + /// without the `Option`. + /// + /// ```no_run + /// use pixelforge::{ + /// ColorConverterConfig, ColorRange, ColorSpec, EncodeConfig, InputFormat, OutputFormat, + /// }; + /// + /// # fn main() -> Result<(), Box> { + /// let conv = ColorConverterConfig::new( + /// 1920, + /// 1080, + /// InputFormat::BGRA, + /// OutputFormat::P010, + /// ColorSpec::Bt709Linear, + /// ColorSpec::Bt2020Pq, + /// ColorRange::Full, + /// ); + /// + /// let enc = EncodeConfig::h265(1920, 1080) + /// .with_color_description(conv.color_description().expect("PQ is encodable")); + /// # Ok(()) + /// # } + /// ``` + pub fn color_description(&self) -> Option { + Some( + self.target + .color_description()? + .with_full_range(self.range.is_full()), + ) + } + + /// Whether the shader can get from this source to this target. + /// + /// Anything that would need tone mapping or a forward gamma encode is not + /// implemented, and is rejected rather than silently passed through. + fn conversion_supported(&self) -> bool { + match self.target { + // Anything can be converted to HDR. + ColorSpec::Bt2020Pq => true, + // Converting to SDR would need tone mapping or a gamma curve. + ColorSpec::Srgb => matches!(self.source, ColorSpec::Srgb), + // Cannot be encoded at all, see `ColorSpec::is_encodable`. + ColorSpec::Bt709Linear | ColorSpec::Bt2020Linear => false, } } + + /// Error explaining why [`Self::conversion_supported`] said no. + fn unsupported_conversion(&self) -> PixelForgeError { + let reason = if !self.target.is_encodable() { + "video cannot be encoded in a linear space, so it can only be a source" + } else { + "converting to SDR would need tone mapping or a gamma curve applied, \ + which the shader does not do" + }; + PixelForgeError::InvalidInput(format!( + "no conversion from {:?} to {:?}: {reason}", + self.source, self.target + )) + } } /// GPU-based color format converter. @@ -241,7 +404,14 @@ pub struct ColorConverter { impl ColorConverter { /// Create a new color converter. + /// + /// Fails if the shader has no path from the configured + /// [`source`](ColorConverterConfig::source) to its + /// [`target`](ColorConverterConfig::target). pub fn new(context: VideoContext, config: ColorConverterConfig) -> Result { + if !config.conversion_supported() { + return Err(config.unsupported_conversion()); + } pipeline::create_converter(context, config) } @@ -252,6 +422,17 @@ impl ColorConverter { &self.context } + /// What to tell the encoder about the colours this converter produces, for + /// [`EncodeConfig::with_color_description`](crate::EncodeConfig::with_color_description). + /// + /// [`ColorConverterConfig::color_description`] without the `Option`: this + /// converter exists, so its target was already checked. + pub fn color_description(&self) -> ColorDescription { + self.config + .color_description() + .expect("a converter cannot be built with a target that has no description") + } + /// Get the converter configuration. pub fn config(&self) -> &ColorConverterConfig { &self.config @@ -265,29 +446,59 @@ impl ColorConverter { self.output_buffer } - /// Set the color space for subsequent conversions. + /// Change what the input pixels are, for subsequent conversions. /// - /// This takes effect on the next `convert()` call without recreating the pipeline, - /// since the color space is passed via push constants. - pub fn set_color_space(&mut self, color_space: ColorSpace) { - self.config.color_space = color_space; + /// Takes effect on the next `convert()` without recreating the pipeline, + /// since the source is passed via push constants. + /// + /// Returns an error, and leaves the converter untouched, if the new source + /// cannot reach the configured target. + pub fn set_source(&mut self, source: ColorSpec) -> Result<()> { + let candidate = ColorConverterConfig { + source, + ..self.config.clone() + }; + if !candidate.conversion_supported() { + return Err(candidate.unsupported_conversion()); + } + self.config.source = source; + Ok(()) } - /// Set the full-range flag for subsequent conversions. + /// Change what the encoded stream should be, for subsequent conversions. /// - /// This takes effect on the next `convert()` call without recreating the pipeline, - /// since the full-range flag is passed via push constants. - pub fn set_full_range(&mut self, full_range: bool) { - self.config.full_range = full_range; + /// Takes effect on the next `convert()` without recreating the pipeline, + /// since the target is passed via push constants. + /// + /// Returns an error, and leaves the converter untouched, if the configured + /// source cannot reach the new target. + /// + /// The target is half of what [`ColorConverterConfig::color_description`] + /// derives from, so an encoder told about the old one is now describing the + /// wrong signal. Pass the new description to + /// [`Encoder::set_color_description`](crate::Encoder::set_color_description), + /// which rebuilds the session parameters and makes the next frame an IDR + /// carrying the updated header. + pub fn set_target(&mut self, target: ColorSpec) -> Result<()> { + let candidate = ColorConverterConfig { + target, + ..self.config.clone() + }; + if !candidate.conversion_supported() { + return Err(candidate.unsupported_conversion()); + } + self.config.target = target; + Ok(()) } - /// Set the SDR reference white level for subsequent conversions. + /// Change the quantization range for subsequent conversions. /// - /// This takes effect on the next `convert()` call without recreating the pipeline, - /// since the value is passed via push constants. Only consumed by the - /// `SrgbToBt2020Pq` and `Bt709LinearToBt2020Pq` color spaces. - pub fn set_sdr_reference_white_nits(&mut self, nits: f32) { - self.config.sdr_reference_white_nits = nits; + /// Takes effect on the next `convert()` without recreating the pipeline, + /// since the range is passed via push constants. The range is the other + /// half of [`ColorConverterConfig::color_description`], so the same + /// re-declaration as [`Self::set_target`] applies. + pub fn set_range(&mut self, range: ColorRange) { + self.config.range = range; } /// Build buffer-to-image copy regions for multi-planar YUV formats. @@ -740,15 +951,17 @@ impl ColorConverter { &offsets, ); - // Push constants: width, height, input_format, output_format, color_space, full_range, sdr_white_nits. - let push_constants: [u32; 7] = [ + // Push constants: width, height, input_format, output_format, + // source, target, range, reference_white_nits. + let push_constants: [u32; 8] = [ self.config.width, self.config.height, self.config.input_format as u32, self.config.output_format as u32, - self.config.color_space as u32, - self.config.full_range as u32, - self.config.sdr_reference_white_nits.to_bits(), + self.config.source as u32, + self.config.target as u32, + self.config.range.is_full() as u32, + self.config.effective_reference_white_nits().to_bits(), ]; let push_constants_bytes: &[u8] = std::slice::from_raw_parts( push_constants.as_ptr() as *const u8, @@ -1010,17 +1223,207 @@ mod tests { } // ======================== - // ColorSpace tests. + // ColorSpec / ColorSpec / ColorRange tests. // ======================== #[test] - fn test_color_space_enum_values() { - // Verify enum values match shader push constant expectations: - // 0=BT.709, 1=BT.2020, 2=sRGB→BT.2020+PQ, 3=scRGB (BT.709 linear)→BT.2020+PQ. - assert_eq!(ColorSpace::Bt709 as u32, 0); - assert_eq!(ColorSpace::Bt2020 as u32, 1); - assert_eq!(ColorSpace::SrgbToBt2020Pq as u32, 2); - assert_eq!(ColorSpace::Bt709LinearToBt2020Pq as u32, 3); + fn spec_discriminants_match_the_shader() { + // The shader's SPEC_* defines. One set, used for both push constants. + assert_eq!(ColorSpec::Srgb as u32, 0); + assert_eq!(ColorSpec::Bt709Linear as u32, 1); + assert_eq!(ColorSpec::Bt2020Linear as u32, 2); + assert_eq!(ColorSpec::Bt2020Pq as u32, 3); + } + + #[test] + fn range_matches_the_shader() { + // The shader's `range == 0u` test. + assert!(!ColorRange::Limited.is_full()); + assert!(ColorRange::Full.is_full()); + } + + #[test] + fn only_encodable_specs_describe_a_stream() { + // A decoder can be told about these two. + assert_eq!( + ColorSpec::Srgb.color_description(), + Some(ColorDescription::bt709()) + ); + assert_eq!( + ColorSpec::Bt2020Pq.color_description(), + Some(ColorDescription::bt2020_pq()) + ); + assert!(ColorSpec::Srgb.is_encodable()); + assert!(ColorSpec::Bt2020Pq.is_encodable()); + + // Linear light cannot be encoded, and these say so. + for spec in [ColorSpec::Bt709Linear, ColorSpec::Bt2020Linear] { + assert!(!spec.is_encodable(), "{spec:?}"); + assert_eq!(spec.color_description(), None, "{spec:?}"); + } + } + + #[test] + fn scrgb_white_is_not_the_srgb_reference() { + // scRGB's white is 80 nits, not the 203 the other spaces use. Using + // 203 for it makes HDR output far too bright. + assert_eq!(ColorSpec::Bt709Linear.reference_white_nits(), Some(80.0)); + assert_eq!(ColorSpec::Srgb.reference_white_nits(), Some(203.0)); + assert_eq!(ColorSpec::Bt2020Linear.reference_white_nits(), Some(203.0)); + } + + #[test] + fn pq_has_no_reference_white() { + // Already carries real brightness values, so there is nothing to ask. + assert_eq!(ColorSpec::Bt2020Pq.reference_white_nits(), None); + } + + #[test] + fn the_override_replaces_the_standard_figure() { + let base = ColorConverterConfig::new( + 64, + 64, + InputFormat::BGRA, + OutputFormat::P010, + ColorSpec::Bt709Linear, + ColorSpec::Bt2020Pq, + ColorRange::Full, + ); + // Unset, the source's own figure is what the shader gets. + assert_eq!(base.reference_white_nits, None); + assert_eq!(base.effective_reference_white_nits(), 80.0); + + let overridden = base.clone().with_reference_white_nits(203.0); + assert_eq!(overridden.effective_reference_white_nits(), 203.0); + // The space itself is untouched; the override lives on the conversion. + assert_eq!(overridden.source, ColorSpec::Bt709Linear); + assert_eq!(overridden.source.reference_white_nits(), Some(80.0)); + } + + #[test] + fn a_pq_source_needs_no_reference_white() { + // Nothing reads it, so an absent figure has to be harmless rather than + // a panic or a NaN reaching the shader. + let config = ColorConverterConfig::new( + 64, + 64, + InputFormat::BGRA, + OutputFormat::P010, + ColorSpec::Bt2020Pq, + ColorSpec::Bt2020Pq, + ColorRange::Full, + ); + assert_eq!(config.effective_reference_white_nits(), 0.0); + } + + // ======================== + // color_description() tests. + // ======================== + + fn config_for(target: ColorSpec, range: ColorRange) -> ColorConverterConfig { + config_from(ColorSpec::Srgb, target, range) + } + + fn config_from( + source: ColorSpec, + target: ColorSpec, + range: ColorRange, + ) -> ColorConverterConfig { + ColorConverterConfig::new( + 64, + 64, + InputFormat::BGRA, + OutputFormat::NV12, + source, + target, + range, + ) + } + + #[test] + fn description_follows_the_target() { + let sdr = config_for(ColorSpec::Srgb, ColorRange::Limited) + .color_description() + .expect("sRGB is encodable"); + assert_eq!(sdr, ColorDescription::bt709()); + assert!(!sdr.is_hdr()); + + let hdr = config_for(ColorSpec::Bt2020Pq, ColorRange::Limited) + .color_description() + .expect("PQ is encodable"); + assert_eq!(hdr, ColorDescription::bt2020_pq()); + assert!(hdr.is_hdr()); + } + + #[test] + fn description_follows_the_range() { + // The drift this whole split exists to prevent: what the shader + // quantizes to and what the VUI claims are now one decision. + for target in [ColorSpec::Srgb, ColorSpec::Bt2020Pq] { + for range in [ColorRange::Limited, ColorRange::Full] { + let config = config_for(target, range); + assert_eq!( + config + .color_description() + .expect("both targets are encodable") + .full_range, + config.range.is_full(), + "{target:?} / {range:?}" + ); + } + } + } + + #[test] + fn description_ignores_the_source() { + // The source says nothing about the encoded stream, so it must not + // reach the declaration. + let from_srgb = config_from(ColorSpec::Srgb, ColorSpec::Bt2020Pq, ColorRange::Full); + let from_scrgb = config_from( + ColorSpec::Bt709Linear, + ColorSpec::Bt2020Pq, + ColorRange::Full, + ); + let from_pq = config_from(ColorSpec::Bt2020Pq, ColorSpec::Bt2020Pq, ColorRange::Full); + assert_eq!( + from_srgb.color_description(), + from_scrgb.color_description() + ); + assert_eq!(from_srgb.color_description(), from_pq.color_description()); + } + + // ======================== + // Supported conversion tests. + // ======================== + + #[test] + fn every_source_can_reach_pq() { + for source in [ + ColorSpec::Srgb, + ColorSpec::Bt709Linear, + ColorSpec::Bt2020Linear, + ColorSpec::Bt2020Pq, + ] { + let config = config_from(source, ColorSpec::Bt2020Pq, ColorRange::Full); + assert!(config.conversion_supported(), "{source:?}"); + } + } + + #[test] + fn only_srgb_can_reach_sdr() { + // Everything else would need a forward gamma encode or tone mapping, + // and the shader does neither. Rejected beats silently passed through. + let ok = config_from(ColorSpec::Srgb, ColorSpec::Srgb, ColorRange::Limited); + assert!(ok.conversion_supported()); + + for source in [ + ColorSpec::Bt709Linear, + ColorSpec::Bt2020Linear, + ColorSpec::Bt2020Pq, + ] { + let config = config_from(source, ColorSpec::Srgb, ColorRange::Limited); + assert!(!config.conversion_supported(), "{source:?}"); + } } // ======================== @@ -1103,36 +1506,36 @@ mod tests { #[test] fn test_config_clone() { - let config = ColorConverterConfig { - width: 1920, - height: 1080, - input_format: InputFormat::BGRx, - output_format: OutputFormat::NV12, - color_space: ColorSpace::Bt709, - full_range: true, - sdr_reference_white_nits: 203.0, - }; + let config = ColorConverterConfig::new( + 1920, + 1080, + InputFormat::BGRx, + OutputFormat::NV12, + ColorSpec::Srgb, + ColorSpec::Srgb, + ColorRange::Full, + ); let cloned = config.clone(); assert_eq!(cloned.width, 1920); assert_eq!(cloned.height, 1080); assert_eq!(cloned.input_format, InputFormat::BGRx); assert_eq!(cloned.output_format, OutputFormat::NV12); - assert_eq!(cloned.color_space, ColorSpace::Bt709); - assert!(cloned.full_range); + assert_eq!(cloned.target, ColorSpec::Srgb); + assert!(cloned.range.is_full()); } #[test] fn test_config_debug() { - let config = ColorConverterConfig { - width: 640, - height: 480, - input_format: InputFormat::RGBA, - output_format: OutputFormat::I420, - color_space: ColorSpace::Bt709, - full_range: true, - sdr_reference_white_nits: 203.0, - }; + let config = ColorConverterConfig::new( + 640, + 480, + InputFormat::RGBA, + OutputFormat::I420, + ColorSpec::Srgb, + ColorSpec::Srgb, + ColorRange::Full, + ); let debug_str = format!("{:?}", config); assert!(debug_str.contains("640")); @@ -1166,15 +1569,15 @@ mod tests { return; }; - let config = ColorConverterConfig { - width: 64, - height: 64, - input_format: InputFormat::BGRx, - output_format: OutputFormat::NV12, - color_space: ColorSpace::Bt709, - full_range: true, - sdr_reference_white_nits: 203.0, - }; + let config = ColorConverterConfig::new( + 64, + 64, + InputFormat::BGRx, + OutputFormat::NV12, + ColorSpec::Srgb, + ColorSpec::Srgb, + ColorRange::Full, + ); let result = ColorConverter::new(context, config); assert!( @@ -1201,15 +1604,15 @@ mod tests { for input_format in &input_formats { for output_format in &output_formats { - let config = ColorConverterConfig { - width: 32, - height: 32, - input_format: *input_format, - output_format: *output_format, - color_space: ColorSpace::Bt709, - full_range: true, - sdr_reference_white_nits: 203.0, - }; + let config = ColorConverterConfig::new( + 32, + 32, + *input_format, + *output_format, + ColorSpec::Srgb, + ColorSpec::Srgb, + ColorRange::Full, + ); let result = ColorConverter::new(context.clone(), config); assert!( diff --git a/src/converter/pipeline.rs b/src/converter/pipeline.rs index fe135dc..24ce281 100644 --- a/src/converter/pipeline.rs +++ b/src/converter/pipeline.rs @@ -16,7 +16,7 @@ const COLOR_CONVERT_SPIRV_BYTES: &[u8] = include_bytes!("../../shader/color_conv /// Get the SPIR-V bytecode for the color conversion shader. /// /// The shader expects: -/// - Push constants: width, height, input_format, output_format, color_space, full_range (6 × u32) +/// - Push constants: width, height, input_format, output_format, source, target, range, reference_white_nits (8 × u32) /// - Binding 0: Input image (sampler2D) /// - Binding 1: Output buffer (YUV data) /// @@ -73,7 +73,7 @@ pub fn create_converter( let push_constant_range = vk::PushConstantRange::default() .stage_flags(vk::ShaderStageFlags::COMPUTE) .offset(0) - .size(28); // 7 x u32: width, height, input_format, output_format, color_space, full_range, sdr_white_nits(f32) + .size(32); // 8 x u32: width, height, input_format, output_format, source, target, range, reference_white_nits(f32) let pipeline_layout_info = vk::PipelineLayoutCreateInfo::default() .set_layouts(std::slice::from_ref(&descriptor_set_layout)) diff --git a/src/lib.rs b/src/lib.rs index f65439b..c6598ee 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -179,31 +179,67 @@ //! //! ## Color Conversion (RGB → YUV) //! -//! PixelForge includes a GPU compute shader for converting RGB input to YUV output, supporting multiple color spaces: -//! -//! | Color Space | Description | -//! |-------------|-------------| -//! | `Bt709` | Standard SDR (BT.709 coefficients) | -//! | `Bt2020` | HDR passthrough (BT.2020 coefficients, PQ-encoded input) | -//! | `SrgbToBt2020Pq` | SDR-in-HDR (sRGB → linear → BT.2020 gamut → PQ OETF) | -//! | `Bt709LinearToBt2020Pq` | scRGB HDR (linear BT.709 → BT.2020 gamut → PQ OETF). `sdr_reference_white_nits` sets the interpretation of 1.0; per the scRGB spec (IEC 61966-2-2), 80 nits. | +//! PixelForge includes a GPU compute shader for converting RGB input to YUV +//! output. The source color describes the input. The target color describes the stream. +//! Both are a [`ColorSpec`]. +//! +//! | `ColorSpec` | | Can be a target | +//! |-------------|-|-----------------| +//! | `Srgb` | Ordinary SDR content | yes | +//! | `Bt709Linear` | scRGB, from an `EXTENDED_SRGB_LINEAR_EXT` swapchain | no | +//! | `Bt2020Linear` | Linear light, wide gamut | no | +//! | `Bt2020Pq` | HDR10, from an `HDR10_ST2084_EXT` swapchain | yes | +//! +//! The linear spaces cannot be a target, because a video file has no way to +//! record that it holds linear light. Any source can be converted to +//! `Bt2020Pq`. Only `Srgb` can be converted to `Srgb`; the others would need +//! tone mapping or a gamma curve applied, which the shader does not do. +//! [`ColorConverter::new`] rejects the combinations it cannot do. +//! +//! Encoding to HDR needs to know how bright the source's white is, since HDR +//! carries real brightness values and SDR does not. Each space has a sensible +//! default, see [`ColorSpec::reference_white_nits`], overridable with +//! [`ColorConverterConfig::with_reference_white_nits`]. //! //! Supported input formats: BGRx, RGBx, BGRA, RGBA, ABGR2101010 (10-bit packed), RGBA16F (FP16). //! Supported output formats: NV12 (8-bit), I420 (8-bit), YUV444 (8-bit), P010 (10-bit), YUV444P10 (10-bit). //! +//! Pass [`ColorConverter::color_description`] to the encoder rather than +//! describing the colours a second time by hand. The two have to agree: if the +//! converter writes full-range pixels and the stream says limited, players +//! stretch the range again and the picture comes out wrong. +//! //! ```rust,no_run -//! use pixelforge::{ColorConverter, ColorConverterConfig, ColorSpace, InputFormat, OutputFormat, VideoContextBuilder}; +//! use pixelforge::{ +//! Codec, ColorConverter, ColorConverterConfig, ColorRange, ColorSpec, EncodeConfig, +//! Encoder, InputFormat, OutputFormat, VideoContextBuilder, +//! }; //! //! # fn main() -> Result<(), Box> { //! let context = VideoContextBuilder::new() //! .app_name("Color Converter") +//! .require_encode(Codec::H265) //! .build()?; //! -//! let mut config = ColorConverterConfig::new(1920, 1080, InputFormat::BGRx, OutputFormat::NV12); -//! config.color_space = ColorSpace::SrgbToBt2020Pq; +//! // SDR desktop content, encoded as HDR10 for an HDR streaming session. +//! let config = ColorConverterConfig::new( +//! 1920, +//! 1080, +//! InputFormat::BGRx, +//! OutputFormat::P010, +//! ColorSpec::Srgb, +//! ColorSpec::Bt2020Pq, +//! ColorRange::Full, +//! ); //! //! let mut converter = ColorConverter::new(context.clone(), config)?; -//! // converter.convert(input_image, output_buffer)?; +//! +//! // The encoder declares exactly what the shader wrote. +//! let encode_config = EncodeConfig::h265(1920, 1080) +//! .with_color_description(converter.color_description()); +//! +//! let mut encoder = Encoder::new(context, encode_config)?; +//! // converter.convert(input_image, layout, encoder.input_image())?; //! # Ok(()) //! # } //! ``` @@ -282,7 +318,9 @@ pub(crate) const fn align4(size: usize) -> usize { (size + 3) & !3 } -pub use converter::{ColorConverter, ColorConverterConfig, ColorSpace, InputFormat, OutputFormat}; +pub use converter::{ + ColorConverter, ColorConverterConfig, ColorRange, ColorSpec, InputFormat, OutputFormat, +}; pub use decoder::{ DecodeConfig, DecodeSink, DecodeSource, DecodeStatus, DecodedFrame, Decoder, FramePoll, Framing, }; diff --git a/tests/color_conversion.rs b/tests/color_conversion.rs new file mode 100644 index 0000000..362a117 --- /dev/null +++ b/tests/color_conversion.rs @@ -0,0 +1,611 @@ +//! What the colour converter writes, checked against the same maths on the CPU. +//! +//! Colour bugs here do not look like failures. They look like a normal picture +//! that is slightly too dark, or the wrong shade. So rather than eyeball it, +//! these tests compute what each conversion should produce and compare. +//! +//! **Quantization.** The shader turns floating-point YUV into integer values. +//! If it truncates instead of rounding, every pixel comes out up to one step +//! too dark. Each sample is scored against both possibilities. +//! +//! **Conversion paths.** Which [`ColorSpec`] pair you pick decides which steps +//! the shader runs: an sRGB decode, a gamut conversion, a PQ encode, or none of +//! them. The same steps are redone on the CPU, picked from the source and +//! target rather than from the shader's branches, so a wrong branch shows up as +//! a large disagreement. +//! +//! Luma only, since chroma is averaged over 2x2 blocks and that averaging would +//! muddy the comparison. The test frame is deliberately not grey: grey converts +//! to exact values and would hide the very errors being looked for. +//! +//! Ignored by default: requires a Vulkan Video device. Run with +//! `cargo test -- --ignored`. + +#[allow(dead_code)] +mod common; + +use ash::vk; +use pixelforge::{ + Codec, ColorConverter, ColorConverterConfig, ColorDescription, ColorRange, ColorSpec, + EncodeBitDepth, EncodeConfig, Encoder, InputFormat, OutputFormat, RateControlMode, + VideoContext, VideoContextBuilder, +}; + +const WIDTH: u32 = 256; +const HEIGHT: u32 = 64; + +/// Deterministic colours that are never grey, so every channel differs and the +/// luma lands between code values often enough to be worth measuring. +fn colour_at(x: u32, y: u32) -> [u8; 3] { + [ + (x ^ y.wrapping_mul(37)) as u8, + x.wrapping_mul(7).wrapping_add(y.wrapping_mul(11)) as u8, + x.wrapping_mul(13).wrapping_add(y.wrapping_mul(3)) as u8, + ] +} + +fn make_frame_bgra() -> Vec { + let mut data = Vec::with_capacity((WIDTH * HEIGHT * 4) as usize); + for y in 0..HEIGHT { + for x in 0..WIDTH { + let [r, g, b] = colour_at(x, y); + data.extend_from_slice(&[b, g, r, 255]); + } + } + data +} + +// --- CPU model of each shader stage, same order, same constants. --- + +fn srgb_to_linear(c: f32) -> f32 { + if c <= 0.04045 { + c / 12.92 + } else { + ((c + 0.055) / 1.055).powf(2.4) + } +} + +fn bt709_to_bt2020(rgb: [f32; 3]) -> [f32; 3] { + let [r, g, b] = rgb; + [ + 0.6274 * r + 0.3293 * g + 0.0433 * b, + 0.0691 * r + 0.9195 * g + 0.0114 * b, + 0.0164 * r + 0.0880 * g + 0.8956 * b, + ] +} + +fn linear_to_pq(l: f32) -> f32 { + // Computed in f64. This is the reference; the shader's own f32 rounding is + // what the one-code-value tolerance covers. + const M1: f64 = 0.1593017578125; + const M2: f64 = 78.84375; + const C1: f64 = 0.8359375; + const C2: f64 = 18.8515625; + const C3: f64 = 18.6875; + let lm1 = (l.max(0.0) as f64).powf(M1); + (((C1 + C2 * lm1) / (1.0 + C3 * lm1)).powf(M2)) as f32 +} + +/// The signal the shader should arrive at for this source and target, then its +/// luma, then its code value. Mirrors `read_rgb` and `rgb_to_yuv`. +fn expected_code( + rgb: [u8; 3], + source: ColorSpec, + target: ColorSpec, + ten_bit: bool, + full: bool, +) -> f32 { + let mut v = [ + rgb[0] as f32 / 255.0, + rgb[1] as f32 / 255.0, + rgb[2] as f32 / 255.0, + ]; + if target == ColorSpec::Bt2020Pq && source != ColorSpec::Bt2020Pq { + if matches!(source, ColorSpec::Srgb) { + v = [ + srgb_to_linear(v[0]), + srgb_to_linear(v[1]), + srgb_to_linear(v[2]), + ]; + } + if !matches!(source, ColorSpec::Bt2020Linear) { + v = bt709_to_bt2020(v); + } + let nits = source.reference_white_nits().unwrap(); + v = [ + linear_to_pq(v[0] * (nits / 10000.0)), + linear_to_pq(v[1] * (nits / 10000.0)), + linear_to_pq(v[2] * (nits / 10000.0)), + ]; + } + let y = if target == ColorSpec::Bt2020Pq { + 0.2627f32 * v[0] + 0.6780f32 * v[1] + 0.0593f32 * v[2] + } else { + 0.2126f32 * v[0] + 0.7152f32 * v[1] + 0.0722f32 * v[2] + } + .clamp(0.0, 1.0); + match (ten_bit, full) { + (false, true) => (y * 255.0).clamp(0.0, 255.0), + (false, false) => (y * 219.0 + 16.0).clamp(0.0, 255.0), + (true, true) => (y * 1023.0).clamp(0.0, 1023.0), + (true, false) => (y * 876.0 + 64.0).clamp(0.0, 1023.0), + } +} + +/// The synthetic frame, on the GPU, cleaned up when the test drops it. +struct SrcImage { + /// Keeps the device alive until after the image is destroyed. + context: VideoContext, + image: vk::Image, + memory: vk::DeviceMemory, +} + +impl Drop for SrcImage { + fn drop(&mut self) { + let device = self.context.device(); + unsafe { + device.destroy_image(self.image, None); + device.free_memory(self.memory, None); + } + } +} + +unsafe fn one_shot( + context: &VideoContext, + record: F, +) -> Result<(), Box> { + let device = context.device(); + let pool_info = vk::CommandPoolCreateInfo::default() + .queue_family_index(context.transfer_queue_family()) + .flags(vk::CommandPoolCreateFlags::TRANSIENT); + let pool = unsafe { device.create_command_pool(&pool_info, None) }?; + let cb_info = vk::CommandBufferAllocateInfo::default() + .command_pool(pool) + .level(vk::CommandBufferLevel::PRIMARY) + .command_buffer_count(1); + let cb = unsafe { device.allocate_command_buffers(&cb_info) }?[0]; + unsafe { + device.begin_command_buffer( + cb, + &vk::CommandBufferBeginInfo::default() + .flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT), + )?; + record(cb); + device.end_command_buffer(cb)?; + let cbs = [cb]; + let submit = vk::SubmitInfo::default().command_buffers(&cbs); + device.queue_submit(context.transfer_queue(), &[submit], vk::Fence::null())?; + device.queue_wait_idle(context.transfer_queue())?; + device.destroy_command_pool(pool, None); + } + Ok(()) +} + +unsafe fn host_buffer( + context: &VideoContext, + size: u64, + usage: vk::BufferUsageFlags, +) -> Result<(vk::Buffer, vk::DeviceMemory), Box> { + let device = context.device(); + let info = vk::BufferCreateInfo::default() + .size(size) + .usage(usage) + .sharing_mode(vk::SharingMode::EXCLUSIVE); + let buffer = unsafe { device.create_buffer(&info, None) }?; + let reqs = unsafe { device.get_buffer_memory_requirements(buffer) }; + let mem_type = context + .find_memory_type( + reqs.memory_type_bits, + vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT, + ) + .ok_or("no host visible memory")?; + let alloc = vk::MemoryAllocateInfo::default() + .allocation_size(reqs.size) + .memory_type_index(mem_type); + let memory = unsafe { device.allocate_memory(&alloc, None) }?; + unsafe { device.bind_buffer_memory(buffer, memory, 0) }?; + Ok((buffer, memory)) +} + +unsafe fn create_src_image(context: &VideoContext) -> Result> { + let device = context.device(); + let info = vk::ImageCreateInfo::default() + .image_type(vk::ImageType::TYPE_2D) + .format(vk::Format::B8G8R8A8_UNORM) + .extent(vk::Extent3D { + width: WIDTH, + height: HEIGHT, + depth: 1, + }) + .mip_levels(1) + .array_layers(1) + .samples(vk::SampleCountFlags::TYPE_1) + .tiling(vk::ImageTiling::OPTIMAL) + .usage(vk::ImageUsageFlags::SAMPLED | vk::ImageUsageFlags::TRANSFER_DST) + .sharing_mode(vk::SharingMode::EXCLUSIVE) + .initial_layout(vk::ImageLayout::UNDEFINED); + let image = unsafe { device.create_image(&info, None) }?; + let reqs = unsafe { device.get_image_memory_requirements(image) }; + let mem_type = context + .find_memory_type(reqs.memory_type_bits, vk::MemoryPropertyFlags::DEVICE_LOCAL) + .ok_or("no device local memory")?; + let alloc = vk::MemoryAllocateInfo::default() + .allocation_size(reqs.size) + .memory_type_index(mem_type); + let memory = unsafe { device.allocate_memory(&alloc, None) }?; + unsafe { device.bind_image_memory(image, memory, 0) }?; + + let pixels = make_frame_bgra(); + let (staging, staging_mem) = unsafe { + host_buffer( + context, + pixels.len() as u64, + vk::BufferUsageFlags::TRANSFER_SRC, + )? + }; + unsafe { + let ptr = device.map_memory( + staging_mem, + 0, + pixels.len() as u64, + vk::MemoryMapFlags::empty(), + )?; + std::ptr::copy_nonoverlapping(pixels.as_ptr(), ptr as *mut u8, pixels.len()); + device.unmap_memory(staging_mem); + } + + let range = vk::ImageSubresourceRange { + aspect_mask: vk::ImageAspectFlags::COLOR, + base_mip_level: 0, + level_count: 1, + base_array_layer: 0, + layer_count: 1, + }; + unsafe { + one_shot(context, |cb| { + let to_dst = vk::ImageMemoryBarrier::default() + .old_layout(vk::ImageLayout::UNDEFINED) + .new_layout(vk::ImageLayout::TRANSFER_DST_OPTIMAL) + .src_queue_family_index(vk::QUEUE_FAMILY_IGNORED) + .dst_queue_family_index(vk::QUEUE_FAMILY_IGNORED) + .image(image) + .subresource_range(range) + .dst_access_mask(vk::AccessFlags::TRANSFER_WRITE); + device.cmd_pipeline_barrier( + cb, + vk::PipelineStageFlags::TOP_OF_PIPE, + vk::PipelineStageFlags::TRANSFER, + vk::DependencyFlags::empty(), + &[], + &[], + &[to_dst], + ); + let copy = vk::BufferImageCopy::default() + .image_subresource(vk::ImageSubresourceLayers { + aspect_mask: vk::ImageAspectFlags::COLOR, + mip_level: 0, + base_array_layer: 0, + layer_count: 1, + }) + .image_extent(vk::Extent3D { + width: WIDTH, + height: HEIGHT, + depth: 1, + }); + device.cmd_copy_buffer_to_image( + cb, + staging, + image, + vk::ImageLayout::TRANSFER_DST_OPTIMAL, + &[copy], + ); + let to_read = vk::ImageMemoryBarrier::default() + .old_layout(vk::ImageLayout::TRANSFER_DST_OPTIMAL) + .new_layout(vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL) + .src_queue_family_index(vk::QUEUE_FAMILY_IGNORED) + .dst_queue_family_index(vk::QUEUE_FAMILY_IGNORED) + .image(image) + .subresource_range(range) + .src_access_mask(vk::AccessFlags::TRANSFER_WRITE) + .dst_access_mask(vk::AccessFlags::SHADER_READ); + device.cmd_pipeline_barrier( + cb, + vk::PipelineStageFlags::TRANSFER, + vk::PipelineStageFlags::ALL_COMMANDS, + vk::DependencyFlags::empty(), + &[], + &[], + &[to_read], + ); + })?; + device.destroy_buffer(staging, None); + device.free_memory(staging_mem, None); + } + Ok(SrcImage { + context: context.clone(), + image, + memory, + }) +} + +/// Read the converter's luma plane back as code values. +fn read_luma( + context: &VideoContext, + converter: &ColorConverter, + output_format: OutputFormat, +) -> Result, Box> { + let sample_bytes = output_format.bytes_per_sample() as u64; + let plane_bytes = (WIDTH * HEIGHT) as u64 * sample_bytes; + let device = context.device(); + let (readback, readback_mem) = + unsafe { host_buffer(context, plane_bytes, vk::BufferUsageFlags::TRANSFER_DST)? }; + unsafe { + one_shot(context, |cb| { + let copy = vk::BufferCopy::default().size(plane_bytes); + device.cmd_copy_buffer(cb, converter.output_buffer(), readback, &[copy]); + })?; + } + let mut bytes = vec![0u8; plane_bytes as usize]; + unsafe { + let ptr = device.map_memory(readback_mem, 0, plane_bytes, vk::MemoryMapFlags::empty())?; + std::ptr::copy_nonoverlapping(ptr as *const u8, bytes.as_mut_ptr(), bytes.len()); + device.unmap_memory(readback_mem); + device.destroy_buffer(readback, None); + device.free_memory(readback_mem, None); + } + Ok(if sample_bytes == 2 { + bytes + .as_chunks::<2>() + .0 + .iter() + .map(|w| (u16::from_le_bytes(*w) >> 6) as u32) + .collect() + } else { + bytes.iter().map(|&b| b as u32).collect() + }) +} + +/// One 8-bit and one 10-bit encoder, for the conversion's target image. +/// +/// 10-bit needs a Main10 encoder, which not every device exposes, so a missing +/// one skips the 10-bit half rather than failing the test. +fn encoders( + context: &VideoContext, +) -> Result<(Encoder, Option), Box> { + let eight = Encoder::new( + context.clone(), + EncodeConfig::h264(WIDTH, HEIGHT) + .with_rate_control(RateControlMode::Cqp) + .with_frame_rate(30, 1) + .with_b_frames(0), + )?; + let ten = Encoder::new( + context.clone(), + EncodeConfig::h265(WIDTH, HEIGHT) + .with_bit_depth(EncodeBitDepth::Ten) + .with_rate_control(RateControlMode::Cqp) + .with_frame_rate(30, 1) + .with_b_frames(0), + ) + .ok(); + if ten.is_none() { + println!("no 10-bit H.265 encoder on this device: P010 cases skipped"); + } + Ok((eight, ten)) +} + +fn context() -> Result> { + Ok(VideoContextBuilder::new() + .app_name("pixelforge-color-conversion") + .require_encode(Codec::H264) + .enable_validation(std::env::var("PIXELFORGE_VALIDATION").is_ok()) + .build()?) +} + +/// Convert one frame and read its luma plane back as code values. +fn convert( + context: &VideoContext, + src: &SrcImage, + encoder: &Encoder, + source: ColorSpec, + target: ColorSpec, + output_format: OutputFormat, + range: ColorRange, +) -> Result<(Vec, ColorDescription), Box> { + let config = ColorConverterConfig::new( + WIDTH, + HEIGHT, + InputFormat::BGRA, + output_format, + source, + target, + range, + ); + let mut converter = ColorConverter::new(context.clone(), config)?; + let description = converter.color_description(); + converter.convert( + src.image, + vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL, + encoder.input_image(), + )?; + Ok((read_luma(context, &converter, output_format)?, description)) +} + +#[test] +#[ignore = "requires a Vulkan Video device"] +fn quantizers_round_rather_than_truncate() -> Result<(), Box> { + let context = context()?; + let src = unsafe { create_src_image(&context)? }; + let (eight, ten) = encoders(&context)?; + + for (output_format, encoder) in [ + (OutputFormat::NV12, Some(&eight)), + (OutputFormat::P010, ten.as_ref()), + ] { + let Some(encoder) = encoder else { continue }; + let ten_bit = output_format.bytes_per_sample() == 2; + for range in [ColorRange::Limited, ColorRange::Full] { + let full = range.is_full(); + let (codes, _) = convert( + &context, + &src, + encoder, + ColorSpec::Srgb, + ColorSpec::Srgb, + output_format, + range, + )?; + + let mut if_truncating = 0usize; + let mut if_rounding = 0usize; + let mut signed_error = 0f64; + for y in 0..HEIGHT { + for x in 0..WIDTH { + let got = codes[(y * WIDTH + x) as usize]; + let ideal = expected_code( + colour_at(x, y), + ColorSpec::Srgb, + ColorSpec::Srgb, + ten_bit, + full, + ); + if got != ideal as u32 { + if_truncating += 1; + } + if got != ideal.round() as u32 { + if_rounding += 1; + } + signed_error += got as f64 - ideal as f64; + } + } + let samples = (WIDTH * HEIGHT) as f64; + let bias = signed_error / samples; + println!( + "{output_format:?} {range:?}: {if_truncating} mismatches if truncating, \ + {if_rounding} if rounding, mean signed error {bias:+.4}" + ); + + // Truncation biases every sample low by about half a code value, + // which is what this is really watching for. The tie tolerance is + // for the GPU ordering its f32 arithmetic differently from us. + assert!( + if_rounding < if_truncating, + "{output_format:?} {range:?}: looks like it truncates" + ); + assert!( + bias.abs() < 0.05, + "{output_format:?} {range:?}: mean signed error {bias:+.4} is a bias, not noise" + ); + } + } + Ok(()) +} + +#[test] +#[ignore = "requires a Vulkan Video device"] +fn every_conversion_matches_the_model() -> Result<(), Box> { + let context = context()?; + let src = unsafe { create_src_image(&context)? }; + let (eight, ten) = encoders(&context)?; + + let supported = [ + (ColorSpec::Srgb, ColorSpec::Srgb), + (ColorSpec::Srgb, ColorSpec::Bt2020Pq), + (ColorSpec::Bt709Linear, ColorSpec::Bt2020Pq), + (ColorSpec::Bt2020Linear, ColorSpec::Bt2020Pq), + (ColorSpec::Bt2020Pq, ColorSpec::Bt2020Pq), + ]; + + for (source, target) in supported { + for (output_format, encoder) in [ + (OutputFormat::NV12, Some(&eight)), + (OutputFormat::P010, ten.as_ref()), + ] { + let Some(encoder) = encoder else { continue }; + let ten_bit = output_format.bytes_per_sample() == 2; + let (codes, description) = convert( + &context, + &src, + encoder, + source, + target, + output_format, + ColorRange::Full, + )?; + + let mut worst = 0f64; + let mut sum = 0f64; + for y in 0..HEIGHT { + for x in 0..WIDTH { + let got = codes[(y * WIDTH + x) as usize] as f64; + let want = expected_code(colour_at(x, y), source, target, ten_bit, true).round() + as f64; + let d = (got - want).abs(); + worst = worst.max(d); + sum += d; + } + } + let mean = sum / (WIDTH * HEIGHT) as f64; + println!( + "{source:?} -> {target:?} {output_format:?}: worst {worst:.1}, mean {mean:.4}, \ + VUI {}/{}/{} {}", + description.color_primaries, + description.transfer_characteristics, + description.matrix_coefficients, + if description.full_range { "pc" } else { "tv" }, + ); + + // The declaration has to describe the target, not the source. + assert_eq!( + Some(description.with_full_range(false)), + target.color_description(), + "{source:?} -> {target:?}: declared the wrong space" + ); + // One code value of slack: the GPU is free to order and fuse these + // f32 operations differently, which moves samples landing on a tie. + assert!( + worst <= 1.0, + "{source:?} -> {target:?} {output_format:?}: worst {worst:.1} code values off" + ); + } + } + Ok(()) +} + +#[test] +#[ignore = "requires a Vulkan Video device"] +fn unsupported_conversions_are_refused() -> Result<(), Box> { + let context = context()?; + + let refused = [ + // Converting to SDR: would need a gamma curve, or tone mapping. + (ColorSpec::Bt709Linear, ColorSpec::Srgb), + (ColorSpec::Bt2020Linear, ColorSpec::Srgb), + (ColorSpec::Bt2020Pq, ColorSpec::Srgb), + // Video cannot be encoded in a linear space. + (ColorSpec::Srgb, ColorSpec::Bt709Linear), + (ColorSpec::Srgb, ColorSpec::Bt2020Linear), + ]; + + for (source, target) in refused { + let config = ColorConverterConfig::new( + WIDTH, + HEIGHT, + InputFormat::BGRA, + OutputFormat::NV12, + source, + target, + ColorRange::Full, + ); + assert_eq!( + config.color_description().is_some(), + target.is_encodable(), + "{target:?}: color_description() and is_encodable() disagree" + ); + let error = ColorConverter::new(context.clone(), config) + .err() + .unwrap_or_else(|| panic!("{source:?} -> {target:?} was accepted")); + println!("{source:?} -> {target:?}: {error}"); + } + Ok(()) +}