From fe9dedd84234b7186113ebf0432d092eed675e4c Mon Sep 17 00:00:00 2001 From: Andreas Reich Date: Tue, 8 Oct 2024 16:58:46 +0200 Subject: [PATCH] chroma subsampling converter now passes all the relevant data now to the converting shader --- .../chroma_subsampling_converter.wgsl | 33 +++++ .../chroma_subsampling_converter.rs | 120 ++++++++++++++---- .../image_data_to_texture.rs | 18 +-- .../re_renderer/src/resource_managers/mod.rs | 2 +- .../src/gpu_bridge/image_to_gpu.rs | 6 +- 5 files changed, 144 insertions(+), 35 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 b704ea8f06b5..251c37f6318d 100644 --- a/crates/viewer/re_renderer/shader/conversions/chroma_subsampling_converter.wgsl +++ b/crates/viewer/re_renderer/shader/conversions/chroma_subsampling_converter.wgsl @@ -1,7 +1,40 @@ #import <../types.wgsl> #import <../screen_triangle_vertex.wgsl> +struct UniformBuffer { + format: u32, + primaries: u32, +}; + +@group(0) @binding(0) +var uniform_buffer: UniformBuffer; + +@group(1) @binding(1) +var input_texture: texture_2d; + + +const FORMAT_Y_UV = 0u; +const FORMAT_YUYV16 = 1u; + +const PRIMARIES_BT601 = 0u; +const PRIMARIES_BT709 = 1u; + + @fragment fn fs_main(in: FragmentInput) -> @location(0) vec4f { + + switch (uniform_buffer.format) { + case FORMAT_Y_UV: { + return vec4f(0.0, 0.0, 1.0, 1.0); + } + case FORMAT_YUYV16: { + return vec4f(1.0, 0.0, 1.0, 1.0); + } + default: { + // Something went wrong! + return vec4f(0.0, 0.0, 0.0, 0.0); + } + } + return vec4f(1.0, 0.0, 1.0, 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 362dbfa00d60..9878de8fed51 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 @@ -1,21 +1,23 @@ use smallvec::smallvec; use crate::{ + allocator::create_and_fill_uniform_buffer, include_shader_module, renderer::{screen_triangle_vertex_shader, DrawData, DrawError, Renderer}, wgpu_resources::{ - GpuBindGroupLayoutHandle, GpuRenderPipelineHandle, GpuTexture, PipelineLayoutDesc, - RenderPipelineDesc, TextureDesc, + BindGroupDesc, BindGroupEntry, BindGroupLayoutDesc, GpuBindGroup, GpuBindGroupLayoutHandle, + GpuRenderPipelineHandle, GpuTexture, PipelineLayoutDesc, RenderPipelineDesc, TextureDesc, }, DebugLabel, RenderContext, }; -use super::ColorSpace; +use super::ColorPrimaries; /// Supported chroma subsampling input formats. /// /// This is strongly correlated with the chroma formats in [`super::SourceImageDataFormat`], but has to have /// stable indices for use in the a uniform buffer. +/// -> Keep in sync with `chroma_subsampling_converter.wgsl` #[allow(non_camel_case_types)] #[derive(Clone, Copy)] pub enum ChromaSubsamplingPixelFormat { @@ -52,20 +54,21 @@ pub enum ChromaSubsamplingPixelFormat { impl ChromaSubsamplingPixelFormat { pub fn expected_data_width_height(&self, width: u32, height: u32) -> (u32, u32) { match self { - ChromaSubsamplingPixelFormat::Y_UV12 => (width, height + height / 2), - ChromaSubsamplingPixelFormat::YUYV16 => (width * 2, height), + Self::Y_UV12 => (width, height + height / 2), + Self::YUYV16 => (width * 2, height), } } pub fn expected_data_texture_format(&self) -> wgpu::TextureFormat { // TODO(andreas): How to deal with higher precision formats here? // If done right that should be easy to drop in. + // TODO(andreas): Are we set on integer formats here, does this actually make sense? #[allow(clippy::match_same_arms)] match self { - ChromaSubsamplingPixelFormat::Y_UV12 => wgpu::TextureFormat::R8Uint, + Self::Y_UV12 => wgpu::TextureFormat::R8Uint, // TODO(andreas): Why not use [`wgpu::TextureFormat::Rg8Uint`] here? - ChromaSubsamplingPixelFormat::YUYV16 => wgpu::TextureFormat::R8Uint, + Self::YUYV16 => wgpu::TextureFormat::R8Uint, } } @@ -73,8 +76,8 @@ impl ChromaSubsamplingPixelFormat { pub fn expected_data_buffer_size(&self, width: u32, height: u32) -> usize { let num_pixels = width as usize * height as usize; match self { - ChromaSubsamplingPixelFormat::Y_UV12 => 12 * num_pixels / 8, - ChromaSubsamplingPixelFormat::YUYV16 => 16 * num_pixels / 8, + Self::Y_UV12 => 12 * num_pixels / 8, + Self::YUYV16 => 16 * num_pixels / 8, } } } @@ -85,18 +88,19 @@ mod gpu_data { #[repr(C)] #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)] pub struct UniformBuffer { - pub format: u32, // Uses [`ChromaSubsamplingPixelFormat`]. + /// Uses [`super::ChromaSubsamplingPixelFormat`]. + pub format: u32, + /// Uses [`super::ColorPrimaries`]. + pub primaries: u32, - pub _padding: [u32; 3], + pub _padding: [u32; 2], pub _end_padding: [wgpu_buffer_types::PaddingRow; 16 - 1], } } -struct ChromaSubsamplingConverterUniformBuffer {} - /// A work item for the subsampling converter. pub struct ChromaSubsamplingConversionTask { - //bind_group: GpuBindGroup, + bind_group: GpuBindGroup, target_texture: GpuTexture, } @@ -118,16 +122,16 @@ impl ChromaSubsamplingConversionTask { pub fn new( ctx: &RenderContext, format: ChromaSubsamplingPixelFormat, - color_space: ColorSpace, - input_data: GpuTexture, - output_label: DebugLabel, + primaries: ColorPrimaries, + input_data: &GpuTexture, + output_label: &DebugLabel, output_width: u32, output_height: u32, ) -> Self { let target_texture = ctx.gpu_resources.textures.alloc( &ctx.device, &TextureDesc { - label: output_label.into(), + label: output_label.clone(), size: wgpu::Extent3d { width: output_width, height: output_height, @@ -143,7 +147,37 @@ impl ChromaSubsamplingConversionTask { }, ); - Self { target_texture } + let renderer = ctx.renderer::(); + + let uniform_buffer = create_and_fill_uniform_buffer( + ctx, + format!("{output_label}_conversion").into(), + gpu_data::UniformBuffer { + format: format as _, + primaries: primaries as _, + + _padding: Default::default(), + _end_padding: Default::default(), + }, + ); + + let bind_group = ctx.gpu_resources.bind_groups.alloc( + &ctx.device, + &ctx.gpu_resources, + &BindGroupDesc { + label: "RectangleInstance::bind_group".into(), + entries: smallvec![ + uniform_buffer, + BindGroupEntry::DefaultTextureView(input_data.handle), + ], + layout: renderer.bind_group_layout, + }, + ); + + Self { + bind_group, + target_texture, + } } /// Runs the conversion from the input texture data. @@ -189,6 +223,7 @@ impl ChromaSubsamplingConversionTask { /// (we need some place to lazily create the render pipeline, store a handle to it and encapsulate the draw logic!) pub struct ChromaSubsamplingConverter { render_pipeline: GpuRenderPipelineHandle, + bind_group_layout: GpuBindGroupLayoutHandle, } impl Renderer for ChromaSubsamplingConverter { @@ -197,11 +232,47 @@ impl Renderer for ChromaSubsamplingConverter { fn create_renderer(ctx: &RenderContext) -> Self { let vertex_handle = screen_triangle_vertex_shader(ctx); + let bind_group_layout = ctx.gpu_resources.bind_group_layouts.get_or_create( + &ctx.device, + &BindGroupLayoutDesc { + label: "ChromaSubsamplingConverter".into(), + entries: vec![ + // Uniform buffer with some information. + wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: wgpu::ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Uniform, + has_dynamic_offset: false, + min_binding_size: (std::mem::size_of::() + as u64) + .try_into() + .ok(), + }, + count: None, + }, + // Input data texture. + wgpu::BindGroupLayoutEntry { + binding: 1, + visibility: wgpu::ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Texture { + multisampled: false, + view_dimension: wgpu::TextureViewDimension::D2, + sample_type: wgpu::TextureSampleType::Uint, + }, + count: None, + }, + ], + }, + ); + let pipeline_layout = ctx.gpu_resources.pipeline_layouts.get_or_create( ctx, &PipelineLayoutDesc { label: "ChromaSubsamplingConverter".into(), - entries: vec![], + // Note that this is a fairly unusual layout for us with the first entry + // not being the globally set bind group! + entries: vec![bind_group_layout], }, ); @@ -230,7 +301,10 @@ impl Renderer for ChromaSubsamplingConverter { }, ); - Self { render_pipeline } + Self { + render_pipeline, + bind_group_layout, + } } fn draw( @@ -238,12 +312,12 @@ impl Renderer for ChromaSubsamplingConverter { render_pipelines: &crate::wgpu_resources::GpuRenderPipelinePoolAccessor<'_>, _phase: crate::draw_phases::DrawPhase, pass: &mut wgpu::RenderPass<'_>, - _draw_data: &Self::RendererDrawData, + draw_data: &Self::RendererDrawData, ) -> Result<(), DrawError> { let pipeline = render_pipelines.get(self.render_pipeline)?; pass.set_pipeline(pipeline); - //pass.set_bind_group(1, &draw_data.bind_group, &[]); + pass.set_bind_group(0, &draw_data.bind_group, &[]); pass.draw(0..3, 0..1); Ok(()) diff --git a/crates/viewer/re_renderer/src/resource_managers/image_data_to_texture.rs b/crates/viewer/re_renderer/src/resource_managers/image_data_to_texture.rs index a55a2c17dd72..0b0198eba3f0 100644 --- a/crates/viewer/re_renderer/src/resource_managers/image_data_to_texture.rs +++ b/crates/viewer/re_renderer/src/resource_managers/image_data_to_texture.rs @@ -7,7 +7,7 @@ use crate::{ DebugLabel, RenderContext, Texture2DBufferInfo, }; -/// Type of color space a given image is in. +/// 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 [`ColorSpace::Bt709`]). @@ -16,8 +16,10 @@ use crate::{ /// /// Ffmpeg's documentation has a short & good overview of these relationships: /// +/// +/// Values need to be kept in sync with `chroma_subsampling_converter.wgsl` #[derive(Clone, Copy, Debug)] -pub enum ColorSpace { +pub enum ColorPrimaries { /// BT.601 (aka. SDTV, aka. Rec.601) /// /// Wiki: @@ -70,12 +72,12 @@ pub enum SourceImageDataFormat { /// /// First comes entire image in Y in one plane, /// followed by a plane with interleaved lines ordered as U0, V0, U1, V1, etc. - Y_UV12(ColorSpace), + Y_UV12(ColorPrimaries), /// `YUYV16` (aka `YUYV` or `YUV2`), is a YUV 4:2:2 chroma downsampled format with 16 bits per pixel and 8 bits per channel. /// /// The order of the channels is Y0, U0, Y1, V0, all in the same plane. - YUYV16(ColorSpace), + YUYV16(ColorPrimaries), } impl From for SourceImageDataFormat { @@ -274,7 +276,7 @@ pub fn transfer_image_data_to_texture( // No further conversion needed, we're done here! return Ok(data_texture); } - SourceImageDataFormat::Y_UV12(color_space) | SourceImageDataFormat::YUYV16(color_space) => { + SourceImageDataFormat::Y_UV12(primaries) | SourceImageDataFormat::YUYV16(primaries) => { let chroma_format = match source_format { SourceImageDataFormat::WgpuCompatible(_) => unreachable!(), SourceImageDataFormat::Y_UV12(_) => ChromaSubsamplingPixelFormat::Y_UV12, @@ -283,9 +285,9 @@ pub fn transfer_image_data_to_texture( ChromaSubsamplingConversionTask::new( ctx, chroma_format, - color_space, - data_texture, - label.clone(), + primaries, + &data_texture, + &label, width, height, ) diff --git a/crates/viewer/re_renderer/src/resource_managers/mod.rs b/crates/viewer/re_renderer/src/resource_managers/mod.rs index 2f2db765daac..0778faca1b9f 100644 --- a/crates/viewer/re_renderer/src/resource_managers/mod.rs +++ b/crates/viewer/re_renderer/src/resource_managers/mod.rs @@ -11,6 +11,6 @@ mod image_data_to_texture; mod texture_manager; pub use image_data_to_texture::{ - ColorSpace, ImageDataDesc, ImageDataToTextureError, SourceImageDataFormat, + ColorPrimaries, ImageDataDesc, ImageDataToTextureError, SourceImageDataFormat, }; pub use texture_manager::{GpuTexture2D, TextureManager2D, TextureManager2DError}; diff --git a/crates/viewer/re_viewer_context/src/gpu_bridge/image_to_gpu.rs b/crates/viewer/re_viewer_context/src/gpu_bridge/image_to_gpu.rs index 3a9c601443cf..12686f3d9dca 100644 --- a/crates/viewer/re_viewer_context/src/gpu_bridge/image_to_gpu.rs +++ b/crates/viewer/re_viewer_context/src/gpu_bridge/image_to_gpu.rs @@ -10,7 +10,7 @@ use re_renderer::{ config::DeviceCaps, pad_rgb_to_rgba, renderer::{ColorMapper, ColormappedTexture, ShaderDecoding}, - resource_managers::{ColorSpace, ImageDataDesc, SourceImageDataFormat}, + resource_managers::{ColorPrimaries, ImageDataDesc, SourceImageDataFormat}, RenderContext, }; use re_types::components::ClassId; @@ -263,11 +263,11 @@ pub fn texture_creation_desc_from_color_image<'a>( // but should confirm & back that up! PixelFormat::NV12 => ( cast_slice_to_cow(image.buffer.as_slice()), - SourceImageDataFormat::Y_UV12(ColorSpace::Bt601), + SourceImageDataFormat::Y_UV12(ColorPrimaries::Bt601), ), PixelFormat::YUY2 => ( cast_slice_to_cow(image.buffer.as_slice()), - SourceImageDataFormat::YUYV16(ColorSpace::Bt601), + SourceImageDataFormat::YUYV16(ColorPrimaries::Bt601), ), } } else {