Skip to content

Commit

Permalink
Introduce image data conversion pipeline, taking over existing YUV co…
Browse files Browse the repository at this point in the history
…nversions (#7640)

### What

* Major part of #7608
* Related to #3931
* since it makes the rectangle shader smaller which kept acting up in
the past


via `pixi run -e py python
./tests/python/chroma_downsample_image/main.py`:

https://rerun.io/viewer/pr/7640?url=https%3A%2F%2Fstatic.rerun.io%2Frrd%2F0.19%2Fchroma_formats_ad9697d05933f4104fbbe66af23073ad4e9e4a58.rrd


![image](https://github.com/user-attachments/assets/76756c5f-75db-44b7-8db0-f04256f72e3a)



The new functionality is tightly integrated into `TextureManager2D`
which now ingests all incoming data via a new internal method
`transfer_image_data_to_texture`.
This in turn takes care of data upload via `GpuReadCpuWriteBelt` (new!
previously, we would use `queue.write_texture` in this case) and may if
needed run gpu based conversion.

Gpu based conversion is for convenience implemented like a `Renderer` -
while it is not _used_ like a typical `Renderer`, the lifecycle (create
lazily once, store on context, feed with data bundles [...]) was so
similar that it made sense to me to use this existing pattern even
though it's not a perfect match.

**A curious sideeffect of this is that you can now put chroma
downsampled textures on custom meshes!**

Next steps:
* support formats required for AV1
   * while on it, check if things could get simpler by us  
* move BGR handling into the new pipeline
* revisit the re_renderer sided input data description. Not sure if I'm
happy about it drifting away from the user facing one!
* once we're there we're also very close to get rid of
`SourceImageDataFormat::WgpuCompatible`. But it's not strictly
necessary, this might yet prove a convenient loophole, although its
presence makes me a little bit nervous (details see code comments) 🤔
* consider exposing color primaries to public API
* the only thing keeping us from it is adding a new enum to the color
format!
    * should it be called "primaries" or "color space" 🤔   

### Checklist
* [x] test on web Mac Chrome/Firefox/Safari
* [x] test on web Window Chrome/Firefox
* [ ] test on web Firefox Chrome/Firefox (@jleibs plz 🥺 )
* [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 the web demo (if applicable):
* Using examples from latest `main` build:
[rerun.io/viewer](https://rerun.io/viewer/pr/7640?manifest_url=https://app.rerun.io/version/main/examples_manifest.json)
* Using full set of examples from `nightly` build:
[rerun.io/viewer](https://rerun.io/viewer/pr/7640?manifest_url=https://app.rerun.io/version/nightly/examples_manifest.json)
* [x] The PR title and labels are set such as to maximize their
usefulness for the next release's CHANGELOG
* [x] If applicable, add a new check to the [release
checklist](https://github.com/rerun-io/rerun/blob/main/tests/python/release_checklist)!
* [x] If have noted any breaking changes to the log API in
`CHANGELOG.md` and the migration guide

- [PR Build Summary](https://build.rerun.io/pr/7640)
- [Recent benchmark results](https://build.rerun.io/graphs/crates.html)
- [Wasm size tracking](https://build.rerun.io/graphs/sizes.html)

To run all checks from `main`, comment on the PR with `@rerun-bot
full-check`.
  • Loading branch information
Wumpf authored Oct 9, 2024
1 parent e952502 commit c3af105
Show file tree
Hide file tree
Showing 21 changed files with 984 additions and 475 deletions.
112 changes: 112 additions & 0 deletions crates/viewer/re_renderer/shader/conversions/yuv_converter.wgsl
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
#import <../types.wgsl>
#import <../screen_triangle_vertex.wgsl>

struct UniformBuffer {
format: u32,
primaries: u32,
target_texture_size: vec2u,
};

@group(0) @binding(0)
var<uniform> uniform_buffer: UniformBuffer;

@group(0) @binding(1)
var input_texture: texture_2d<u32>;


const YUV_LAYOUT_Y_UV = 0u;
const YUV_LAYOUT_YUYV16 = 1u;

const PRIMARIES_BT601 = 0u;
const PRIMARIES_BT709 = 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 {
// 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 = ERROR_RGBA.rgb;
}
}

return clamp(rgb, vec3f(0.0), vec3f(1.0));
}

/// Extracts YUV data from a chroma subsampling encoded texture at specific coordinates.
///
/// 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 {
let texture_dim = vec2f(textureDimensions(texture).xy);
var yuv: vec3f;

switch (yuv_layout) {
case YUV_LAYOUT_Y_UV: {
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 YUV_LAYOUT_YUYV16: {
// 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: {
yuv = vec3f(0.0, 0.0, 0.0); // ERROR_RGBA doesn't apply here.
}
}

return yuv;
}

@fragment
fn fs_main(in: FragmentInput) -> @location(0) vec4f {
let coords = 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);

return vec4f(rgb, 1.0);
}
70 changes: 0 additions & 70 deletions crates/viewer/re_renderer/shader/decodings.wgsl

This file was deleted.

17 changes: 0 additions & 17 deletions crates/viewer/re_renderer/shader/rectangle_fs.wgsl
Original file line number Diff line number Diff line change
@@ -1,11 +1,6 @@
#import <./colormap.wgsl>
#import <./rectangle.wgsl>
#import <./utils/srgb.wgsl>
#import <./decodings.wgsl>

// WARNING! Adding anything else to this shader is very likely to push us over a size threshold that
// causes the failure reported in: https://github.com/rerun-io/rerun/issues/3931
// Make sure any changes are tested in Chrome on Linux using the Intel Mesa driver.

fn is_magnifying(pixel_coord: vec2f) -> bool {
return fwidth(pixel_coord.x) < 1.0;
Expand Down Expand Up @@ -86,10 +81,6 @@ fn fs_main(in: VertexOut) -> @location(0) vec4f {
texture_dimensions = vec2f(textureDimensions(texture_sint).xy);
} else if rect_info.sample_type == SAMPLE_TYPE_UINT {
texture_dimensions = vec2f(textureDimensions(texture_uint).xy);
} else if rect_info.sample_type == SAMPLE_TYPE_NV12 {
texture_dimensions = vec2f(textureDimensions(texture_uint).xy);
} else if rect_info.sample_type == SAMPLE_TYPE_YUY2 {
texture_dimensions = vec2f(textureDimensions(texture_uint).xy);
}

let coord = in.texcoord * texture_dimensions;
Expand Down Expand Up @@ -141,14 +132,6 @@ fn fs_main(in: VertexOut) -> @location(0) vec4f {
vec4f(textureLoad(texture_uint, v01_coord, 0)),
vec4f(textureLoad(texture_uint, v10_coord, 0)),
vec4f(textureLoad(texture_uint, v11_coord, 0)));
} else if rect_info.sample_type == SAMPLE_TYPE_NV12 || rect_info.sample_type == SAMPLE_TYPE_YUY2{
normalized_value = decode_color_and_filter_nearest_or_bilinear(
filter_nearest,
coord,
decode_nv12_or_yuy2(rect_info.sample_type, texture_uint, v00_coord),
decode_nv12_or_yuy2(rect_info.sample_type, texture_uint, v01_coord),
decode_nv12_or_yuy2(rect_info.sample_type, texture_uint, v10_coord),
decode_nv12_or_yuy2(rect_info.sample_type, texture_uint, v11_coord));
} else {
return ERROR_RGBA; // unknown sample type
}
Expand Down
3 changes: 1 addition & 2 deletions crates/viewer/re_renderer/src/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -217,8 +217,7 @@ impl RenderContext {
}

let resolver = crate::new_recommended_file_resolver();
let texture_manager_2d =
TextureManager2D::new(device.clone(), queue.clone(), &gpu_resources.textures);
let texture_manager_2d = TextureManager2D::new(&device, &queue, &gpu_resources.textures);

let active_frame = ActiveFrameContext {
before_view_builder_encoder: Mutex::new(FrameGlobalCommandEncoder::new(&device)),
Expand Down
28 changes: 11 additions & 17 deletions crates/viewer/re_renderer/src/importer/gltf.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ use smallvec::SmallVec;
use crate::{
mesh::{GpuMesh, Material, Mesh, MeshError},
renderer::MeshInstance,
resource_managers::{GpuTexture2D, Texture2DCreationDesc, TextureManager2D},
resource_managers::{GpuTexture2D, ImageDataDesc, TextureManager2D},
RenderContext, Rgba32Unmul,
};

Expand Down Expand Up @@ -84,31 +84,25 @@ pub fn load_gltf_from_buffer(
#[cfg(not(debug_assertions))]
let texture_names = "";

let texture = Texture2DCreationDesc {
let texture = ImageDataDesc {
label: if texture_names.is_empty() {
format!("unnamed gltf image in {mesh_name}")
} else {
format!("gltf image used by {texture_names} in {mesh_name}")
}
.into(),
data: data.into(),
format,
width: image.width,
height: image.height,
format: format.into(),
width_height: [image.width, image.height],
};

images_as_textures.push(
match ctx
.texture_manager_2d
.create(&ctx.gpu_resources.textures, &texture)
{
Ok(texture) => texture,
Err(err) => {
re_log::error!("Failed to create texture: {err}");
ctx.texture_manager_2d.white_texture_unorm_handle().clone()
}
},
);
images_as_textures.push(match ctx.texture_manager_2d.create(ctx, texture) {
Ok(texture) => texture,
Err(err) => {
re_log::error!("Failed to create texture: {err}");
ctx.texture_manager_2d.white_texture_unorm_handle().clone()
}
});
}

let mut meshes = HashMap::with_capacity(doc.meshes().len());
Expand Down
2 changes: 2 additions & 0 deletions crates/viewer/re_renderer/src/renderer/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,8 @@ pub trait Renderer {
}

/// Gets or creates a vertex shader module for drawing a screen filling triangle.
///
/// The entry point of this shader is `main`.
pub fn screen_triangle_vertex_shader(
ctx: &RenderContext,
) -> crate::wgpu_resources::GpuShaderModuleHandle {
Expand Down
28 changes: 3 additions & 25 deletions crates/viewer/re_renderer/src/renderer/rectangles.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,11 +47,9 @@ pub enum TextureFilterMin {
}

/// Describes how the color information is encoded in the texture.
// TODO(#7608): to be replaced by re_renderer based on-input conversion.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ShaderDecoding {
Nv12,
Yuy2,

/// Do BGR(A)->RGB(A) conversion is in the shader.
Bgr,
}
Expand Down Expand Up @@ -145,17 +143,7 @@ impl ColormappedTexture {
}

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]
}
Some(ShaderDecoding::Yuy2) => {
let [width, height] = self.texture.width_height();
[width / 2, height]
}
Some(ShaderDecoding::Bgr) | None => self.texture.width_height(),
}
self.texture.width_height()
}
}

Expand Down Expand Up @@ -237,8 +225,6 @@ 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;
const SAMPLE_TYPE_YUY2: u32 = 5;

// How do we do colormapping?
const COLOR_MAPPER_OFF_GRAYSCALE: u32 = 1;
Expand Down Expand Up @@ -318,15 +304,7 @@ mod gpu_data {
let sample_type = match texture_format.sample_type(None, None) {
Some(wgpu::TextureSampleType::Float { .. }) => SAMPLE_TYPE_FLOAT,
Some(wgpu::TextureSampleType::Sint) => SAMPLE_TYPE_SINT,
Some(wgpu::TextureSampleType::Uint) => {
if shader_decoding == &Some(super::ShaderDecoding::Nv12) {
SAMPLE_TYPE_NV12
} else if shader_decoding == &Some(super::ShaderDecoding::Yuy2) {
SAMPLE_TYPE_YUY2
} else {
SAMPLE_TYPE_UINT
}
}
Some(wgpu::TextureSampleType::Uint) => SAMPLE_TYPE_UINT,
_ => {
return Err(RectangleError::TextureFormatNotSupported(texture_format));
}
Expand Down
Loading

0 comments on commit c3af105

Please sign in to comment.