handle all these new foramts and test them in chroma_downsample_image…

…. lotsa fixes etc. on the way

crates/store/re_types/src/datatypes/pixel_format_ext.rs

@@ -1,4 +1,4 @@
-use crate::image::rgb_from_yuv;
+use crate::image::{rgb_from_yuv, ColorPrimaries};
 
 use super::{ChannelDatatype, ColorModel, PixelFormat};
 
@@ -92,33 +92,34 @@ impl PixelFormat {
             Self::Y_U_V24_FullRange | Self::Y_U_V24_LimitedRange => {
                 let plane_size = (w * h) as usize;
                 let plane_coord = (y * w + x) as usize;
+
                 let luma = *buf.get(plane_coord)?;
-                let u = *buf.get(plane_size + plane_coord)?;
-                let v = *buf.get(plane_size * 2 + plane_coord)?;
+                let u = *buf.get(plane_coord + plane_size)?;
+                let v = *buf.get(plane_coord + plane_size * 2)?;
                 Some([luma, u, v])
             }
 
             Self::Y_U_V16_FullRange | Self::Y_U_V16_LimitedRange => {
                 let y_plane_size = (w * h) as usize;
-                let uv_plane_size = y_plane_size / 2;
+                let uv_plane_size = y_plane_size / 2; // Half horizontal resolution.
                 let y_plane_coord = (y * w + x) as usize;
-                let uv_plane_coord = (y * w + x / 2) as usize;
+                let uv_plane_coord = y_plane_coord / 2; // == (y * (w / 2) + x / 2)
 
                 let luma = *buf.get(y_plane_coord)?;
-                let u = *buf.get(y_plane_coord + uv_plane_coord)?;
-                let v = *buf.get(y_plane_coord + uv_plane_size + uv_plane_coord)?;
+                let u = *buf.get(uv_plane_coord + y_plane_size)?;
+                let v = *buf.get(uv_plane_coord + y_plane_size + uv_plane_size)?;
                 Some([luma, u, v])
             }
 
             Self::Y_U_V12_FullRange | Self::Y_U_V12_LimitedRange => {
                 let y_plane_size = (w * h) as usize;
-                let uv_plane_size = y_plane_size / 4;
+                let uv_plane_size = y_plane_size / 4; // Half horizontal & vertical resolution.
                 let y_plane_coord = (y * w + x) as usize;
-                let uv_plane_coord = y_plane_coord / 2;
+                let uv_plane_coord = (y * w / 4 + x / 2) as usize; // == ((y / 2) * (w / 2) + x / 2)
 
                 let luma = *buf.get(y_plane_coord)?;
-                let u = *buf.get(y_plane_coord + uv_plane_coord)?;
-                let v = *buf.get(y_plane_coord + uv_plane_size + uv_plane_coord)?;
+                let u = *buf.get(uv_plane_coord + y_plane_size)?;
+                let v = *buf.get(uv_plane_coord + y_plane_size + uv_plane_size)?;
                 Some([luma, u, v])
             }
 
@@ -142,18 +143,33 @@ impl PixelFormat {
     }
 
     /// Returns true if the format is a YUV format using
-    /// limited range YUV, i.e. Y is valid in [16, 235] and U/V [16, 240],
-    /// rather than 0-255 or higher ranges.
+    /// limited range YUV.
+    ///
+    /// I.e. for 8bit data, Y is valid in [16, 235] and U/V [16, 240], rather than 0-255.
     pub fn is_limited_yuv_range(&self) -> bool {
         match self {
-            Self::Y_U_V12_LimitedRange => true,
-            Self::NV12 => true,
-            Self::YUY2 => true,
-            Self::Y_U_V24_LimitedRange => true,
-            Self::Y_U_V24_FullRange => false,
-            Self::Y_U_V12_FullRange => false,
-            Self::Y_U_V16_LimitedRange => true,
-            Self::Y_U_V16_FullRange => false,
+            Self::Y_U_V24_LimitedRange
+            | Self::Y_U_V16_LimitedRange
+            | Self::Y_U_V12_LimitedRange
+            | Self::NV12
+            | Self::YUY2 => true,
+
+            Self::Y_U_V24_FullRange | Self::Y_U_V12_FullRange | Self::Y_U_V16_FullRange => false,
+        }
+    }
+
+    /// Color primaries used by this format.
+    // TODO(andreas): Expose this in the API separately and document it better.
+    pub fn color_primaries(&self) -> ColorPrimaries {
+        match self {
+            Self::Y_U_V24_LimitedRange
+            | Self::Y_U_V24_FullRange
+            | Self::Y_U_V12_LimitedRange
+            | Self::Y_U_V12_FullRange
+            | Self::Y_U_V16_LimitedRange
+            | Self::Y_U_V16_FullRange => ColorPrimaries::Bt709,
+
+            Self::NV12 | Self::YUY2 => ColorPrimaries::Bt601,
         }
     }
 
@@ -163,6 +179,12 @@ impl PixelFormat {
     #[inline]
     pub fn decode_rgb_at(&self, buf: &[u8], [w, h]: [u32; 2], [x, y]: [u32; 2]) -> Option<[u8; 3]> {
         let [y, u, v] = self.decode_yuv_at(buf, [w, h], [x, y])?;
-        Some(rgb_from_yuv(y, u, v, self.is_limited_yuv_range()))
+        Some(rgb_from_yuv(
+            y,
+            u,
+            v,
+            self.is_limited_yuv_range(),
+            self.color_primaries(),
+        ))
     }
 }

crates/store/re_types/src/image.rs

@@ -266,30 +266,91 @@ fn test_find_non_empty_dim_indices() {
 
 // ----------------------------------------------------------------------------
 
+// TODO(andreas): Expose this in the API.
+/// Type of color primaries a given image is in.
+///
+/// This applies both to YUV and RGB formats, but if not specified otherwise
+/// we assume BT.709 primaries for all RGB(A) 8bits per channel content (details below on [`ColorPrimaries::Bt709`]).
+/// Since with YUV content the color space is often less clear, we always explicitly
+/// specify it.
+///
+/// Ffmpeg's documentation has a short & good overview of these relationships:
+/// <https://trac.ffmpeg.org/wiki/colorspace#WhatiscolorspaceWhyshouldwecare/>
+#[derive(Clone, Copy, Debug)]
+pub enum ColorPrimaries {
+    /// BT.601 (aka. SDTV, aka. Rec.601)
+    ///
+    /// Wiki: <https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion/>
+    Bt601,
+
+    /// BT.709 (aka. HDTV, aka. Rec.709)
+    ///
+    /// Wiki: <https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.709_conversion/>
+    ///
+    /// These are the same primaries we usually assume and use for all our rendering
+    /// since they are the same primaries used by sRGB.
+    /// <https://en.wikipedia.org/wiki/Rec._709#Relationship_to_sRGB/>
+    /// The OETF/EOTF function (<https://en.wikipedia.org/wiki/Transfer_functions_in_imaging/>) is different,
+    /// but for all other purposes they are the same.
+    /// (The only reason for us to convert to optical units ("linear" instead of "gamma") is for
+    /// lighting & tonemapping where we typically start out with an sRGB image!)
+    Bt709,
+    //
+    // Not yet supported. These vary a lot more from the other two!
+    //
+    // /// BT.2020 (aka. PQ, aka. Rec.2020)
+    // ///
+    // /// Wiki: <https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.2020_conversion/>
+    // BT2020_ConstantLuminance,
+    // BT2020_NonConstantLuminance,
+}
+
 /// Returns sRGB from YUV color.
 ///
 /// This conversion mirrors the function of the same name in `yuv_converter.wgsl`
 ///
 /// Specifying the color standard should be exposed in the future [#3541](https://github.com/rerun-io/rerun/pull/3541)
-pub fn rgb_from_yuv(y: u8, u: u8, v: u8, limited_range: bool) -> [u8; 3] {
+pub fn rgb_from_yuv(
+    y: u8,
+    u: u8,
+    v: u8,
+    limited_range: bool,
+    primaries: ColorPrimaries,
+) -> [u8; 3] {
     let (mut y, mut u, mut v) = (y as f32, u as f32, v as f32);
 
     // rescale YUV values
     if limited_range {
+        // "The resultant signals range from 16 to 235 for Y′ (Cb and Cr range from 16 to 240);
+        // the values from 0 to 15 are called footroom, while the values from 236 to 255 are called headroom."
         y = (y - 16.0) / 219.0;
         u = (u - 128.0) / 224.0;
         v = (v - 128.0) / 224.0;
+    } else {
+        y /= 255.0;
+        u = (u - 128.0) / 255.0;
+        v = (v - 128.0) / 255.0;
     }
 
-    // BT.601 (aka. SDTV, aka. Rec.601). wiki: https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion
-    let r = y + 1.402 * v;
-    let g = y - 0.344 * u - 0.714 * v;
-    let b = y + 1.772 * u;
-
-    // BT.709 (aka. HDTV, aka. Rec.709). wiki: https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.709_conversion
-    // let r = y + 1.575 * v;
-    // let g = y - 0.187 * u - 0.468 * v;
-    // let b = y + 1.856 * u;
+    let r;
+    let g;
+    let b;
+
+    match primaries {
+        ColorPrimaries::Bt601 => {
+            // BT.601 (aka. SDTV, aka. Rec.601). wiki: https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion
+            r = y + 1.402 * v;
+            g = y - 0.344 * u - 0.714 * v;
+            b = y + 1.772 * u;
+        }
+
+        ColorPrimaries::Bt709 => {
+            // BT.709 (aka. HDTV, aka. Rec.709). wiki: https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.709_conversion
+            r = y + 1.575 * v;
+            g = y - 0.187 * u - 0.468 * v;
+            b = y + 1.856 * u;
+        }
+    }
 
     [(255.0 * r) as u8, (255.0 * g) as u8, (255.0 * b) as u8]
 }

crates/viewer/re_renderer/shader/conversions/yuv_converter.wgsl

@@ -2,9 +2,10 @@
 #import <../screen_triangle_vertex.wgsl>
 
 struct UniformBuffer {
-    format: u32,
+    yuv_layout: u32,
     primaries: u32,
     target_texture_size: vec2u,
+    yuv_range: u32,
 };
 
 @group(0) @binding(0)
@@ -25,20 +26,44 @@ const YUV_LAYOUT_Y_400 = 300u;
 const PRIMARIES_BT601 = 0u;
 const PRIMARIES_BT709 = 1u;
 
+// see `enum YuvRange`.
+const YUV_RANGE_LIMITED = 0u;
+const YUV_RANGE_FULL = 1u;
+
 
 /// 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 {
+fn srgb_from_yuv(yuv: vec3f, primaries: u32, range: u32) -> vec3f {
     // rescale YUV values
     //
     // This is what is called "limited range" and is the most common case.
     // TODO(andreas): Support "full range" as well.
-    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 y: f32;
+    var u: f32;
+    var v: f32;
+
+    switch (range) {
+        case YUV_RANGE_LIMITED: {
+            y = (yuv[0] - 16.0) / 219.0;
+            u = (yuv[1] - 128.0) / 224.0;
+            v = (yuv[2] - 128.0) / 224.0;
+        }
+
+        case YUV_RANGE_FULL: {
+            y = yuv[0] / 255.0;
+            u = (yuv[1] - 128.0) / 255.0;
+            v = (yuv[2] - 128.0) / 255.0;
+        }
+
+        default: {
+            // Should never happen.
+            return ERROR_RGBA.rgb;
+        }
+    }
 
     var rgb: vec3f;
 
@@ -61,7 +86,7 @@ fn srgb_from_yuv(yuv: vec3f, primaries: u32) -> vec3f {
         }
 
         default: {
-            rgb = ERROR_RGBA.rgb;
+            return ERROR_RGBA.rgb;
         }
     }
 
@@ -72,27 +97,70 @@ fn srgb_from_yuv(yuv: vec3f, primaries: u32) -> vec3f {
 ///
 /// See also `enum YuvPixelLayout` in `yuv_converter.rs for a specification of
 /// the expected data layout.
-fn sample_yuv(yuv_layout: u32, texture: texture_2d<u32>, coords: vec2f) -> vec3f {
+fn sample_yuv(yuv_layout: u32, texture: texture_2d<u32>, coords: vec2u, target_texture_size: vec2u) -> vec3f {
     let texture_dim = vec2f(textureDimensions(texture).xy);
     var yuv: vec3f;
 
     switch (yuv_layout)  {
+        case YUV_LAYOUT_Y_U_V444: {
+            // Just 3 planes under each other.
+            yuv[0] = f32(textureLoad(texture, coords, 0).r);
+            yuv[1] = f32(textureLoad(texture, vec2u(coords.x, coords.y + target_texture_size.y), 0).r);
+            yuv[2] = f32(textureLoad(texture, vec2u(coords.x, coords.y + target_texture_size.y * 2u), 0).r);
+        }
+
+        case YUV_LAYOUT_Y_U_V422: {
+            // A large Y plane, followed by a UV plane with half the horizontal resolution,
+            // every row contains two u/v rows.
+            yuv[0] = f32(textureLoad(texture, coords, 0).r);
+            // UV coordinate on its own plane:
+            let uv_coord = vec2u(coords.x / 2u, coords.y);
+            // UV coordinate on the data texture, ignoring offset from previous planes.
+            // Each texture row contains two UV rows
+            let uv_col = uv_coord.x + (uv_coord.y % 2) * target_texture_size.x / 2u;
+            let uv_row = uv_coord.y / 2u;
+
+            yuv[1] = f32(textureLoad(texture, vec2u(uv_col, uv_row + target_texture_size.y), 0).r);
+            yuv[2] = f32(textureLoad(texture, vec2u(uv_col, uv_row + target_texture_size.y + target_texture_size.y / 2u), 0).r);
+        }
+
+        case YUV_LAYOUT_Y_U_V420: {
+            // A large Y plane, followed by a UV plane with half the horizontal & vertical resolution,
+            // every row contains two u/v rows and there's only half as many.
+            yuv[0] = f32(textureLoad(texture, coords, 0).r);
+            // UV coordinate on its own plane:
+            let uv_coord = vec2u(coords.x / 2u, coords.y / 2u);
+            // UV coordinate on the data texture, ignoring offset from previous planes.
+            // Each texture row contains two UV rows
+            let uv_col = uv_coord.x + (uv_coord.y % 2) * (target_texture_size.x / 2u);
+            let uv_row = uv_coord.y / 2u;
+
+            yuv[1] = f32(textureLoad(texture, vec2u(uv_col, uv_row + target_texture_size.y), 0).r);
+            yuv[2] = f32(textureLoad(texture, vec2u(uv_col, uv_row + target_texture_size.y + target_texture_size.y / 4u), 0).r);
+        }
+
+        case YUV_LAYOUT_Y_400 {
+            yuv[0] = f32(textureLoad(texture, coords, 0).r);
+            yuv[1] = 128.0;
+            yuv[0] = 128.0;
+        }
+
         case YUV_LAYOUT_Y_UV420: {
             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;
+            let uv_row = (coords.y / 2u);
+            var uv_col = (coords.x / 2u) * 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);
+            yuv[0] = f32(textureLoad(texture, coords, 0).r);
+            yuv[1] = f32(textureLoad(texture, vec2u(uv_col, uv_offset + uv_row), 0).r);
+            yuv[2] = f32(textureLoad(texture, vec2u((uv_col + 1u), uv_offset + uv_row), 0).r);
         }
 
         case YUV_LAYOUT_YUYV422: {
             // texture is 2 * width * height
             // every 4 bytes is 2 pixels
-            let uv_row = u32(coords.y);
+            let uv_row = coords.y;
             // multiply by 2 because the width is multiplied by 2
-            let y_col = u32(coords.x) * 2u;
+            let y_col = 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
@@ -111,10 +179,10 @@ fn sample_yuv(yuv_layout: u32, texture: texture_2d<u32>, coords: vec2f) -> vec3f
 
 @fragment
 fn fs_main(in: FragmentInput) -> @location(0) vec4f {
-    let coords = vec2f(uniform_buffer.target_texture_size) * in.texcoord;
+    let coords = vec2u(vec2f(uniform_buffer.target_texture_size) * in.texcoord);
 
-    let yuv = sample_yuv(uniform_buffer.format, input_texture, coords);
-    let rgb = srgb_from_yuv(yuv, uniform_buffer.primaries);
+    let yuv = sample_yuv(uniform_buffer.yuv_layout, input_texture, coords, uniform_buffer.target_texture_size);
+    let rgb = srgb_from_yuv(yuv, uniform_buffer.primaries, uniform_buffer.yuv_range);
 
     return vec4f(rgb, 1.0);
 }

crates/viewer/re_renderer/src/resource_managers/image_data_to_texture.rs

@@ -34,7 +34,7 @@ pub enum ColorPrimaries {
     /// but for all other purposes they are the same.
     /// (The only reason for us to convert to optical units ("linear" instead of "gamma") is for
     /// lighting & tonemapping where we typically start out with an sRGB image!)
-    Bt709 = 2,
+    Bt709 = 1,
     //
     // Not yet supported. These vary a lot more from the other two!
     //