From a853a2e6e23cbc78cc9a6b8c3a703a8e5480dea9 Mon Sep 17 00:00:00 2001 From: Andreas Reich Date: Tue, 8 Oct 2024 17:30:34 +0200 Subject: [PATCH] chroma conversion shader has now parity with rectangle shader chroma decode --- .../chroma_subsampling_converter.wgsl | 90 +++++++++++++++++-- .../chroma_subsampling_converter.rs | 6 +- 2 files changed, 85 insertions(+), 11 deletions(-) diff --git a/crates/viewer/re_renderer/shader/conversions/chroma_subsampling_converter.wgsl b/crates/viewer/re_renderer/shader/conversions/chroma_subsampling_converter.wgsl index 251c37f6318d..a372068c2b2b 100644 --- a/crates/viewer/re_renderer/shader/conversions/chroma_subsampling_converter.wgsl +++ b/crates/viewer/re_renderer/shader/conversions/chroma_subsampling_converter.wgsl @@ -4,12 +4,13 @@ struct UniformBuffer { format: u32, primaries: u32, + target_texture_size: vec2u, }; @group(0) @binding(0) var uniform_buffer: UniformBuffer; -@group(1) @binding(1) +@group(0) @binding(1) var input_texture: texture_2d; @@ -20,21 +21,92 @@ const PRIMARIES_BT601 = 0u; const PRIMARIES_BT709 = 1u; -@fragment -fn fs_main(in: FragmentInput) -> @location(0) vec4f { +/// Returns sRGB from YUV color. +/// +/// This conversion mirrors the function in `crates/store/re_types/src/datatypes/tensor_data_ext.rs` +/// +/// Specifying the color standard should be exposed in the future [#3541](https://github.com/rerun-io/rerun/pull/3541) +fn srgb_from_yuv(yuv: vec3f, primaries: u32) -> vec3f { + // rescale YUV values + let y = (yuv[0] - 16.0) / 219.0; + let u = (yuv[1] - 128.0) / 224.0; + let v = (yuv[2] - 128.0) / 224.0; + + var rgb: vec3f; + + switch (primaries) { + // BT.601 (aka. SDTV, aka. Rec.601). wiki: https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion + // Also note according to https://en.wikipedia.org/wiki/SRGB#sYCC_extended-gamut_transformation + // > Although the RGB color primaries are based on BT.709, + // > the equations for transformation from sRGB to sYCC and vice versa are based on BT.601. + case PRIMARIES_BT601: { + rgb.r = y + 1.402 * v; + rgb.g = y - 0.344 * u - 0.714 * v; + rgb.b = y + 1.772 * u; + } + + // BT.709 (aka. HDTV, aka. Rec.709). wiki: https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.709_conversion + case PRIMARIES_BT709: { + rgb.r = y + 1.575 * v; + rgb.g = y - 0.187 * u - 0.468 * v; + rgb.b = y + 1.856 * u; + } + + default: { + rgb = vec3f(1.0, 0.0, 1.0); // Magenta to indicate trouble + } + } + + return clamp(rgb, vec3f(0.0), vec3f(1.0)); +} + +/// Extracts YUV data from a chroma subsampling encoded texture at specific coordinates. +/// +/// See also `enum ChromaSubsamplingPixelFormat` in `chroma_subsampling_converter.rs for a specification of +/// the expected data layout. +fn decode_chroma_subsampling_format_to_yuv(format: u32, texture: texture_2d, coords: vec2f) -> vec3f { + let texture_dim = vec2f(textureDimensions(texture).xy); + var yuv: vec3f; - switch (uniform_buffer.format) { + switch (format) { case FORMAT_Y_UV: { - return vec4f(0.0, 0.0, 1.0, 1.0); + let uv_offset = u32(floor(texture_dim.y / 1.5)); + let uv_row = u32(coords.y / 2); + var uv_col = u32(coords.x / 2) * 2u; + + yuv[0] = f32(textureLoad(texture, vec2u(coords), 0).r); + yuv[1] = f32(textureLoad(texture, vec2u(u32(uv_col), uv_offset + uv_row), 0).r); + yuv[2] = f32(textureLoad(texture, vec2u((u32(uv_col) + 1u), uv_offset + uv_row), 0).r); } + case FORMAT_YUYV16: { - return vec4f(1.0, 0.0, 1.0, 1.0); + // texture is 2 * width * height + // every 4 bytes is 2 pixels + let uv_row = u32(coords.y); + // multiply by 2 because the width is multiplied by 2 + let y_col = u32(coords.x) * 2u; + yuv[0] = f32(textureLoad(texture, vec2u(y_col, uv_row), 0).r); + + // at odd pixels we're in the second half of the yuyu block, offset back by 2 + let uv_col = y_col - u32(coords.x % 2) * 2u; + yuv[1] = f32(textureLoad(texture, vec2u(uv_col + 1u, uv_row), 0).r); + yuv[2] = f32(textureLoad(texture, vec2u(uv_col + 3u, uv_row), 0).r); } + default: { - // Something went wrong! - return vec4f(0.0, 0.0, 0.0, 0.0); + yuv = vec3f(0.0, 0.0, 0.0); } } - return vec4f(1.0, 0.0, 1.0, 1.0); + return yuv; +} + +@fragment +fn fs_main(in: FragmentInput) -> @location(0) vec4f { + let coords = vec2f(uniform_buffer.target_texture_size) * in.texcoord; + + let yuv = decode_chroma_subsampling_format_to_yuv(uniform_buffer.format, input_texture, coords); + let rgb = srgb_from_yuv(yuv, uniform_buffer.primaries); + + return vec4f(rgb, 1.0); } diff --git a/crates/viewer/re_renderer/src/resource_managers/chroma_subsampling_converter.rs b/crates/viewer/re_renderer/src/resource_managers/chroma_subsampling_converter.rs index 9878de8fed51..13f9a075226b 100644 --- a/crates/viewer/re_renderer/src/resource_managers/chroma_subsampling_converter.rs +++ b/crates/viewer/re_renderer/src/resource_managers/chroma_subsampling_converter.rs @@ -90,10 +90,12 @@ mod gpu_data { pub struct UniformBuffer { /// Uses [`super::ChromaSubsamplingPixelFormat`]. pub format: u32, + /// Uses [`super::ColorPrimaries`]. pub primaries: u32, - pub _padding: [u32; 2], + pub target_texture_size: [u32; 2], + pub _end_padding: [wgpu_buffer_types::PaddingRow; 16 - 1], } } @@ -155,8 +157,8 @@ impl ChromaSubsamplingConversionTask { gpu_data::UniformBuffer { format: format as _, primaries: primaries as _, + target_texture_size: [output_width, output_height], - _padding: Default::default(), _end_padding: Default::default(), }, );