Support NV12-encoded images (#3541)

### Added support for logging NV12 encoded images via the existing log_image_file api. As promissed a while back :) ```py from rerun import ImageFormat import rerun as rr rr.log_image_file("NV12 image", img_bytes=bytes(frame), img_format=ImageFormat.NV12(width=1920, height=1080)) ``` ![NV12 image example](https://github.com/rerun-io/rerun/assets/59307111/de8ef517-d84f-4130-8946-4ea787f0181e) The raw (encoded) image data is stored in an R8Uint texture. The contents of the texture get decoded in the `rectangle_fs.wgsl` via the decoder written in `decodings.wgsl`. Other [YUV](https://gist.github.com/Jim-Bar/3cbba684a71d1a9d468a6711a6eddbeb) image formats could be added easily, following the NV12 implementation with slight modifications to the decoder. ### Checklist * [x] I have read and agree to [Contributor Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and the [Code of Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md) * [x] I've included a screenshot or gif (if applicable) * [x] I have tested [demo.rerun.io](https://demo.rerun.io/pr/3541) (if applicable) - [PR Build Summary](https://build.rerun.io/pr/3541) - [Docs preview](https://rerun.io/preview/fc13dd7aa9ae329b3e5100f8bb70c32f2d1b8add/docs)  - [Examples preview](https://rerun.io/preview/fc13dd7aa9ae329b3e5100f8bb70c32f2d1b8add/examples)  - [Recent benchmark results](https://ref.rerun.io/dev/bench/) - [Wasm size tracking](https://ref.rerun.io/dev/sizes/) --------- Co-authored-by: Andreas Reich <[email protected]>
rerun-io · Oct 16, 2023 · a363a6c · a363a6c
1 parent 11dee21
commit a363a6c
Show file tree

Hide file tree

Showing 34 changed files with 790 additions and 112 deletions.
diff --git a/crates/re_data_ui/src/image.rs b/crates/re_data_ui/src/image.rs
@@ -138,10 +138,18 @@ fn tensor_ui(
                     );
                 }
 
+                let shape = match tensor.image_height_width_channels() {
+                    Some([h, w, c]) => vec![
+                        TensorDimension::height(h),
+                        TensorDimension::width(w),
+                        TensorDimension::depth(c),
+                    ],
+                    None => tensor.shape.clone(),
+                };
                 ui.label(format!(
                     "{} x {}{}",
                     tensor.dtype(),
-                    format_tensor_shape_single_line(tensor.shape()),
+                    format_tensor_shape_single_line(shape.as_slice()),
                     if original_tensor.buffer.is_compressed_image() {
                         " (compressed)"
                     } else {
@@ -216,6 +224,9 @@ fn tensor_ui(
                     }
 
                     if let Some([_h, _w, channels]) = tensor.image_height_width_channels() {
+                        if let TensorBuffer::Nv12(_) = &tensor.buffer {
+                            return;
+                        }
                         if channels == 3 {
                             if let TensorBuffer::U8(data) = &tensor.buffer {
                                 ui.collapsing("Histogram", |ui| {
@@ -231,7 +242,7 @@ fn tensor_ui(
 }
 
 fn texture_size(colormapped_texture: &ColormappedTexture) -> Vec2 {
-    let [w, h] = colormapped_texture.texture.width_height();
+    let [w, h] = colormapped_texture.width_height();
     egui::vec2(w as f32, h as f32)
 }
 
@@ -360,6 +371,11 @@ pub fn tensor_summary_ui_grid_contents(
             ));
             ui.end_row();
         }
+        TensorBuffer::Nv12(_) => {
+            re_ui.grid_left_hand_label(ui, "Encoding");
+            ui.label("NV12");
+            ui.end_row();
+        }
     }
 
     let TensorStats {
@@ -379,8 +395,9 @@ pub fn tensor_summary_ui_grid_contents(
     }
     // Show finite range only if it is different from the actual range.
     if let (true, Some((min, max))) = (range != finite_range, finite_range) {
-        ui.label("Finite data range")
-            .on_hover_text("The finite values (ignoring all NaN & -Inf/+Inf) of the tensor range within these bounds");
+        ui.label("Finite data range").on_hover_text(
+            "The finite values (ignoring all NaN & -Inf/+Inf) of the tensor range within these bounds"
+        );
         ui.monospace(format!(
             "[{} - {}]",
             re_format::format_f64(*min),
@@ -439,8 +456,8 @@ fn show_zoomed_image_region_tooltip(
                     use egui::remap_clamp;
 
                     let center_texel = [
-                        (remap_clamp(pointer_pos.x, image_rect.x_range(), 0.0..=w as f32) as isize),
-                        (remap_clamp(pointer_pos.y, image_rect.y_range(), 0.0..=h as f32) as isize),
+                        remap_clamp(pointer_pos.x, image_rect.x_range(), 0.0..=w as f32) as isize,
+                        remap_clamp(pointer_pos.y, image_rect.y_range(), 0.0..=h as f32) as isize,
                     ];
                     show_zoomed_image_region_area_outline(
                         parent_ui.ctx(),
@@ -562,7 +579,7 @@ fn try_show_zoomed_image_region(
     )?;
 
     const POINTS_PER_TEXEL: f32 = 5.0;
-    let size = Vec2::splat((ZOOMED_IMAGE_TEXEL_RADIUS * 2 + 1) as f32 * POINTS_PER_TEXEL);
+    let size = Vec2::splat(((ZOOMED_IMAGE_TEXEL_RADIUS * 2 + 1) as f32) * POINTS_PER_TEXEL);
 
     let (_id, zoom_rect) = ui.allocate_space(size);
     let painter = ui.painter();
@@ -574,7 +591,10 @@ fn try_show_zoomed_image_region(
         let image_rect_on_screen = egui::Rect::from_min_size(
             zoom_rect.center()
                 - POINTS_PER_TEXEL
-                    * egui::vec2(center_texel[0] as f32 + 0.5, center_texel[1] as f32 + 0.5),
+                    * egui::vec2(
+                        (center_texel[0] as f32) + 0.5,
+                        (center_texel[1] as f32) + 0.5,
+                    ),
             POINTS_PER_TEXEL * egui::vec2(width as f32, height as f32),
         );
 
@@ -610,7 +630,11 @@ fn try_show_zoomed_image_region(
             let zoom = rect.width();
             let image_rect_on_screen = egui::Rect::from_min_size(
                 rect.center()
-                    - zoom * egui::vec2(center_texel[0] as f32 + 0.5, center_texel[1] as f32 + 0.5),
+                    - zoom
+                        * egui::vec2(
+                            (center_texel[0] as f32) + 0.5,
+                            (center_texel[1] as f32) + 0.5,
+                        ),
                 zoom * egui::vec2(width as f32, height as f32),
             );
             gpu_bridge::render_image(
@@ -661,7 +685,7 @@ fn tensor_pixel_value_ui(
             // This is a depth map
             if let Some(raw_value) = tensor.get(&[y, x]) {
                 let raw_value = raw_value.as_f64();
-                let meters = raw_value / meter as f64;
+                let meters = raw_value / (meter as f64);
                 ui.label("Depth:");
                 if meters < 1.0 {
                     ui.monospace(format!("{:.1} mm", meters * 1e3));
@@ -679,11 +703,20 @@ fn tensor_pixel_value_ui(
                 .map(|v| format!("Val: {v}")),
             3 => {
                 // TODO(jleibs): Track RGB ordering somehow -- don't just assume it
-                if let (Some(r), Some(g), Some(b)) = (
-                    tensor.get_with_image_coords(x, y, 0),
-                    tensor.get_with_image_coords(x, y, 1),
-                    tensor.get_with_image_coords(x, y, 2),
-                ) {
+                if let Some([r, g, b]) = match &tensor.buffer {
+                    TensorBuffer::Nv12(_) => tensor.get_nv12_pixel(x, y),
+                    _ => {
+                        if let [Some(r), Some(g), Some(b)] = [
+                            tensor.get_with_image_coords(x, y, 0),
+                            tensor.get_with_image_coords(x, y, 1),
+                            tensor.get_with_image_coords(x, y, 2),
+                        ] {
+                            Some([r, g, b])
+                        } else {
+                            None
+                        }
+                    }
+                } {
                     match (r, g, b) {
                         (TensorElement::U8(r), TensorElement::U8(g), TensorElement::U8(b)) => {
                             Some(format!("R: {r}, G: {g}, B: {b}, #{r:02X}{g:02X}{b:02X}"))

diff --git a/crates/re_renderer/shader/decodings.wgsl b/crates/re_renderer/shader/decodings.wgsl
@@ -0,0 +1,25 @@
+#import <./types.wgsl>
+
+
+/// Loads an RGBA texel from a texture holding an NV12 encoded image at the given screen space coordinates.
+fn decode_nv12(texture: texture_2d<u32>, coords: IVec2) -> Vec4 {
+    let texture_dim = Vec2(textureDimensions(texture).xy);
+    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 y = max(0.0, (f32(textureLoad(texture, UVec2(coords), 0).r) - 16.0)) / 219.0;
+    let u = (f32(textureLoad(texture, UVec2(u32(uv_col), uv_offset + uv_row), 0).r) - 128.0) / 224.0;
+    let v = (f32(textureLoad(texture, UVec2((u32(uv_col) + 1u), uv_offset + uv_row), 0).r) - 128.0) / 224.0;
+
+    // Specifying the color standard should be exposed in the future (https://github.com/rerun-io/rerun/pull/3541)
+    // BT.601 (aka. SDTV, aka. Rec.601). wiki: https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion
+    let r = clamp(y + 1.402 * v, 0.0, 1.0);
+    let g = clamp(y  - (0.344 * u + 0.714 * v), 0.0, 1.0);
+    let b = clamp(y + 1.772 * u, 0.0, 1.0);
+    // BT.709 (aka. HDTV, aka. Rec.709). wiki: https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.709_conversion
+    // let r = clamp(y + 1.5748 * v, 0.0, 1.0);
+    // let g = clamp(y + u * -0.1873 + v * -0.4681, 0.0, 1.0);
+    // let b = clamp(y + u * 1.8556, 0.0 , 1.0);
+    return Vec4(r, g, b, 1.0);
+}
diff --git a/crates/re_renderer/shader/rectangle.wgsl b/crates/re_renderer/shader/rectangle.wgsl
@@ -6,6 +6,7 @@
 const SAMPLE_TYPE_FLOAT = 1u;
 const SAMPLE_TYPE_SINT  = 2u;
 const SAMPLE_TYPE_UINT  = 3u;
+const SAMPLE_TYPE_NV12 = 4u;
 
 // How do we do colormapping?
 const COLOR_MAPPER_OFF      = 1u;

diff --git a/crates/re_renderer/shader/rectangle_fs.wgsl b/crates/re_renderer/shader/rectangle_fs.wgsl
@@ -1,6 +1,7 @@
 #import <./colormap.wgsl>
 #import <./rectangle.wgsl>
 #import <./utils/srgb.wgsl>
+#import <./decodings.wgsl>
 
 fn is_magnifying(pixel_coord: Vec2) -> bool {
     return fwidth(pixel_coord.x) < 1.0;
@@ -101,7 +102,24 @@ fn fs_main(in: VertexOut) -> @location(0) Vec4 {
             let v11 = decode_color(Vec4(textureLoad(texture_uint, clamp_to_edge_nearest_neighbor(coord + vec2( 0.5,  0.5), texture_dimensions), 0)));
             normalized_value = filter_bilinear(coord, v00, v01, v10, v11);
         }
-    } else {
+    } else if rect_info.sample_type == SAMPLE_TYPE_NV12 {
+        let texture_dimensions = Vec2(textureDimensions(texture_uint).xy);
+        let coord = in.texcoord * texture_dimensions;
+        if tex_filter(coord) == FILTER_NEAREST {
+            // nearest
+            normalized_value = decode_color(Vec4(decode_nv12(texture_uint,
+                clamp_to_edge_nearest_neighbor(coord, texture_dimensions))));
+        } else {
+            // bilinear
+            let v00 = decode_color(Vec4(decode_nv12(texture_uint, clamp_to_edge_nearest_neighbor(coord + vec2(-0.5, -0.5), texture_dimensions))));
+            let v01 = decode_color(Vec4(decode_nv12(texture_uint, clamp_to_edge_nearest_neighbor(coord + vec2(-0.5,  0.5), texture_dimensions))));
+            let v10 = decode_color(Vec4(decode_nv12(texture_uint, clamp_to_edge_nearest_neighbor(coord + vec2( 0.5, -0.5), texture_dimensions))));
+            let v11 = decode_color(Vec4(decode_nv12(texture_uint, clamp_to_edge_nearest_neighbor(coord + vec2( 0.5,  0.5), texture_dimensions))));
+            normalized_value = filter_bilinear(coord, v00, v01, v10, v11);
+        }
+    }
+
+    else {
         return ERROR_RGBA; // unknown sample type
     }
 

diff --git a/crates/re_renderer/shader/rectangle_vs.wgsl b/crates/re_renderer/shader/rectangle_vs.wgsl
@@ -10,6 +10,9 @@ fn vs_main(@builtin(vertex_index) v_idx: u32) -> VertexOut {
     var out: VertexOut;
     out.position = apply_depth_offset(frame.projection_from_world * Vec4(pos, 1.0), rect_info.depth_offset);
     out.texcoord = texcoord;
+    if rect_info.sample_type == SAMPLE_TYPE_NV12 {
+        out.texcoord.y /= 1.5;
+    }
 
     return out;
 }
diff --git a/crates/re_renderer/src/renderer/mod.rs b/crates/re_renderer/src/renderer/mod.rs
@@ -21,8 +21,8 @@ pub use test_triangle::TestTriangleDrawData;
 
 mod rectangles;
 pub use rectangles::{
-    ColorMapper, ColormappedTexture, RectangleDrawData, RectangleOptions, TextureFilterMag,
-    TextureFilterMin, TexturedRect,
+    ColorMapper, ColormappedTexture, RectangleDrawData, RectangleOptions, ShaderDecoding,
+    TextureFilterMag, TextureFilterMin, TexturedRect,
 };
 
 mod mesh_renderer;

diff --git a/crates/re_renderer/src/renderer/rectangles.rs b/crates/re_renderer/src/renderer/rectangles.rs
@@ -48,6 +48,12 @@ pub enum TextureFilterMin {
     // TODO(andreas): Offer mipmapping here?
 }
 
+/// Describes how the color information is encoded in the texture.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum ShaderDecoding {
+    Nv12,
+}
+
 /// Describes a texture and how to map it to a color.
 #[derive(Clone)]
 pub struct ColormappedTexture {
@@ -83,6 +89,9 @@ pub struct ColormappedTexture {
     /// Setting a color mapper for a four-component texture is an error.
     /// Failure to set a color mapper for a one-component texture is an error.
     pub color_mapper: Option<ColorMapper>,
+
+    /// For textures that need decoding in the shader, for example NV12 encoded images.
+    pub shader_decoding: Option<ShaderDecoding>,
 }
 
 /// How to map the normalized `.r` component to a color.
@@ -113,6 +122,17 @@ impl ColormappedTexture {
             gamma: 1.0,
             multiply_rgb_with_alpha: true,
             color_mapper: None,
+            shader_decoding: None,
+        }
+    }
+
+    pub fn width_height(&self) -> [u32; 2] {
+        match self.shader_decoding {
+            Some(ShaderDecoding::Nv12) => {
+                let [width, height] = self.texture.width_height();
+                [width, height * 2 / 3]
+            }
+            _ => self.texture.width_height(),
         }
     }
 }
@@ -198,6 +218,7 @@ mod gpu_data {
     const SAMPLE_TYPE_FLOAT: u32 = 1;
     const SAMPLE_TYPE_SINT: u32 = 2;
     const SAMPLE_TYPE_UINT: u32 = 3;
+    const SAMPLE_TYPE_NV12: u32 = 4;
 
     // How do we do colormapping?
     const COLOR_MAPPER_OFF: u32 = 1;
@@ -261,6 +282,7 @@ mod gpu_data {
                 gamma,
                 color_mapper,
                 multiply_rgb_with_alpha,
+                shader_decoding,
             } = colormapped_texture;
 
             let super::RectangleOptions {
@@ -274,7 +296,13 @@ mod gpu_data {
             let sample_type = match texture_format.sample_type(None) {
                 Some(wgpu::TextureSampleType::Float { .. }) => SAMPLE_TYPE_FLOAT,
                 Some(wgpu::TextureSampleType::Sint) => SAMPLE_TYPE_SINT,
-                Some(wgpu::TextureSampleType::Uint) => SAMPLE_TYPE_UINT,
+                Some(wgpu::TextureSampleType::Uint) => {
+                    if shader_decoding == &Some(super::ShaderDecoding::Nv12) {
+                        SAMPLE_TYPE_NV12
+                    } else {
+                        SAMPLE_TYPE_UINT
+                    }
+                }
                 _ => {
                     return Err(RectangleError::TextureFormatNotSupported(texture_format));
                 }
@@ -292,9 +320,10 @@ mod gpu_data {
                     Some(ColorMapper::Texture(_)) => {
                         color_mapper_int = COLOR_MAPPER_TEXTURE;
                     }
-                    None => {
-                        return Err(RectangleError::MissingColorMapper);
-                    }
+                    None => match shader_decoding {
+                        Some(super::ShaderDecoding::Nv12) => color_mapper_int = COLOR_MAPPER_OFF,
+                        _ => return Err(RectangleError::MissingColorMapper),
+                    },
                 },
                 4 => {
                     if color_mapper.is_some() {
@@ -304,7 +333,7 @@ mod gpu_data {
                     }
                 }
                 num_components => {
-                    return Err(RectangleError::UnsupportedComponentCount(num_components))
+                    return Err(RectangleError::UnsupportedComponentCount(num_components));
                 }
             }
 
@@ -442,7 +471,7 @@ impl RectangleDrawData {
                             BindGroupEntry::DefaultTextureView(texture_float),
                             BindGroupEntry::DefaultTextureView(texture_sint),
                             BindGroupEntry::DefaultTextureView(texture_uint),
-                            BindGroupEntry::DefaultTextureView(colormap_texture),
+                            BindGroupEntry::DefaultTextureView(colormap_texture)
                         ],
                         layout: rectangle_renderer.bind_group_layout,
                     },
@@ -475,7 +504,7 @@ impl Renderer for RectangleRenderer {
 
         let bind_group_layout = pools.bind_group_layouts.get_or_create(
             device,
-            &BindGroupLayoutDesc {
+            &(BindGroupLayoutDesc {
                 label: "RectangleRenderer::bind_group_layout".into(),
                 entries: vec![
                     wgpu::BindGroupLayoutEntry {
@@ -538,15 +567,15 @@ impl Renderer for RectangleRenderer {
                         count: None,
                     },
                 ],
-            },
+            }),
         );
 
         let pipeline_layout = pools.pipeline_layouts.get_or_create(
             device,
-            &PipelineLayoutDesc {
+            &(PipelineLayoutDesc {
                 label: "RectangleRenderer::pipeline_layout".into(),
                 entries: vec![shared_data.global_bindings.layout, bind_group_layout],
-            },
+            }),
             &pools.bind_group_layouts,
         );
 
@@ -591,20 +620,20 @@ impl Renderer for RectangleRenderer {
         );
         let render_pipeline_picking_layer = pools.render_pipelines.get_or_create(
             device,
-            &RenderPipelineDesc {
+            &(RenderPipelineDesc {
                 label: "RectangleRenderer::render_pipeline_picking_layer".into(),
                 fragment_entrypoint: "fs_main_picking_layer".into(),
                 render_targets: smallvec![Some(PickingLayerProcessor::PICKING_LAYER_FORMAT.into())],
                 depth_stencil: PickingLayerProcessor::PICKING_LAYER_DEPTH_STATE,
                 multisample: PickingLayerProcessor::PICKING_LAYER_MSAA_STATE,
                 ..render_pipeline_desc_color.clone()
-            },
+            }),
             &pools.pipeline_layouts,
             &pools.shader_modules,
         );
         let render_pipeline_outline_mask = pools.render_pipelines.get_or_create(
             device,
-            &RenderPipelineDesc {
+            &(RenderPipelineDesc {
                 label: "RectangleRenderer::render_pipeline_outline_mask".into(),
                 fragment_entrypoint: "fs_main_outline_mask".into(),
                 render_targets: smallvec![Some(OutlineMaskProcessor::MASK_FORMAT.into())],
@@ -613,7 +642,7 @@ impl Renderer for RectangleRenderer {
                     &shared_data.config.device_caps,
                 ),
                 ..render_pipeline_desc_color
-            },
+            }),
             &pools.pipeline_layouts,
             &pools.shader_modules,
         );