diff options
Diffstat (limited to 'src/demosaic/lmmse.rs')
-rw-r--r-- | src/demosaic/lmmse.rs | 323 |
1 files changed, 170 insertions, 153 deletions
diff --git a/src/demosaic/lmmse.rs b/src/demosaic/lmmse.rs index 1dc3094..9add4c1 100644 --- a/src/demosaic/lmmse.rs +++ b/src/demosaic/lmmse.rs @@ -1,165 +1,182 @@ pub struct Lmmse { - interpolate_bind_layout: wgpu::BindGroupLayout, - interpolate_pipeline_layout: wgpu::PipelineLayout, - interpolate_shader: wgpu::ShaderModule, - interpolate_pipeline: wgpu::ComputePipeline, + interpolate_bind_layout: wgpu::BindGroupLayout, + interpolate_pipeline_layout: wgpu::PipelineLayout, + interpolate_shader: wgpu::ShaderModule, + interpolate_pipeline: wgpu::ComputePipeline, } impl super::Demosaic for Lmmse { - fn new(gpu: &wgpu::Device, queue: &wgpu::Queue) -> Self { - let interpolate_bind_layout = gpu.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: None, - entries: &[ - wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: wgpu::ShaderStages::COMPUTE, - ty: wgpu::BindingType::StorageTexture { - access: wgpu::StorageTextureAccess::ReadOnly, - format: wgpu::TextureFormat::R16Uint, - view_dimension: wgpu::TextureViewDimension::D2, - }, - count: None, - }, - wgpu::BindGroupLayoutEntry { - binding: 1, - visibility: wgpu::ShaderStages::COMPUTE, - ty: wgpu::BindingType::StorageTexture { - access: wgpu::StorageTextureAccess::WriteOnly, - format: wgpu::TextureFormat::Rgba32Float, - view_dimension: wgpu::TextureViewDimension::D2, - }, - count: None, - } - ], - }); - let interpolate_pipeline_layout = gpu.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { - label: None, - bind_group_layouts: &[&interpolate_bind_layout], - push_constant_ranges: &[], - }); - let interpolate_shader = gpu.create_shader_module(wgpu::ShaderModuleDescriptor { - label: None, - source: wgpu::ShaderSource::Wgsl(include_str!("lmmse_01_interpolate.wgsl").into()), - }); - let interpolate_pipeline = gpu.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: None, - layout: Some(&interpolate_pipeline_layout), - module: &interpolate_shader, - entry_point: None, - compilation_options: wgpu::PipelineCompilationOptions::default(), - cache: None, - }); + fn new(gpu: &wgpu::Device, queue: &wgpu::Queue) -> Self { + let interpolate_bind_layout = + gpu.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: None, + entries: &[ + wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::StorageTexture { + access: wgpu::StorageTextureAccess::ReadOnly, + format: wgpu::TextureFormat::R16Uint, + view_dimension: wgpu::TextureViewDimension::D2, + }, + count: None, + }, + wgpu::BindGroupLayoutEntry { + binding: 1, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::StorageTexture { + access: wgpu::StorageTextureAccess::WriteOnly, + format: wgpu::TextureFormat::Rgba32Float, + view_dimension: wgpu::TextureViewDimension::D2, + }, + count: None, + }, + ], + }); + let interpolate_pipeline_layout = + gpu.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: None, + bind_group_layouts: &[&interpolate_bind_layout], + push_constant_ranges: &[], + }); + let interpolate_shader = gpu.create_shader_module(wgpu::ShaderModuleDescriptor { + label: None, + source: wgpu::ShaderSource::Wgsl(include_str!("lmmse_01_interpolate.wgsl").into()), + }); + let interpolate_pipeline = gpu.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: None, + layout: Some(&interpolate_pipeline_layout), + module: &interpolate_shader, + entry_point: None, + compilation_options: wgpu::PipelineCompilationOptions::default(), + cache: None, + }); - Self { - interpolate_bind_layout, - interpolate_pipeline_layout, - interpolate_shader, - interpolate_pipeline, - } - } + Self { + interpolate_bind_layout, + interpolate_pipeline_layout, + interpolate_shader, + interpolate_pipeline, + } + } - fn demoasic(&self, gpu: &wgpu::Device, queue: &wgpu::Queue, image: &rawloader::RawImage) -> wgpu::Texture { - let input_texture = gpu.create_texture(&wgpu::TextureDescriptor { - label: None, - size: wgpu::Extent3d { - width: image.width as u32, - height: image.height as u32, - depth_or_array_layers: 1, - }, - mip_level_count: 1, - sample_count: 1, - dimension: wgpu::TextureDimension::D2, - format: wgpu::TextureFormat::R16Uint, - usage: wgpu::TextureUsages::COPY_DST | wgpu::TextureUsages::STORAGE_BINDING, - view_formats: &[wgpu::TextureFormat::R16Uint], - }); + fn demoasic( + &self, + gpu: &wgpu::Device, + queue: &wgpu::Queue, + image: &rawloader::RawImage, + ) -> wgpu::Texture { + let input_texture = gpu.create_texture(&wgpu::TextureDescriptor { + label: None, + size: wgpu::Extent3d { + width: image.width as u32, + height: image.height as u32, + depth_or_array_layers: 1, + }, + mip_level_count: 1, + sample_count: 1, + dimension: wgpu::TextureDimension::D2, + format: wgpu::TextureFormat::R16Uint, + usage: wgpu::TextureUsages::COPY_DST | wgpu::TextureUsages::STORAGE_BINDING, + view_formats: &[wgpu::TextureFormat::R16Uint], + }); - let image_data = match &image.data { - rawloader::RawImageData::Integer(d) => d.as_slice(), - _ => panic!("Unsupported input format"), - }; - queue.write_texture( - wgpu::TexelCopyTextureInfo { - texture: &input_texture, - mip_level: 0, - origin: wgpu::Origin3d::ZERO, - aspect: wgpu::TextureAspect::All, - }, - bytemuck::cast_slice(image_data), - wgpu::TexelCopyBufferLayout { - offset: 0, - bytes_per_row: Some(2 * image.width as u32), - rows_per_image: Some(image.height as u32), - }, - wgpu::Extent3d { - width: image.width as u32, - height: image.height as u32, - depth_or_array_layers: 1, - } - ); + let image_data = match &image.data { + rawloader::RawImageData::Integer(d) => d.as_slice(), + _ => panic!("Unsupported input format"), + }; + queue.write_texture( + wgpu::TexelCopyTextureInfo { + texture: &input_texture, + mip_level: 0, + origin: wgpu::Origin3d::ZERO, + aspect: wgpu::TextureAspect::All, + }, + bytemuck::cast_slice(image_data), + wgpu::TexelCopyBufferLayout { + offset: 0, + bytes_per_row: Some(2 * image.width as u32), + rows_per_image: Some(image.height as u32), + }, + wgpu::Extent3d { + width: image.width as u32, + height: image.height as u32, + depth_or_array_layers: 1, + }, + ); - let output_texture = gpu.create_texture(&wgpu::TextureDescriptor { - label: None, - size: wgpu::Extent3d { - width: image.width as u32, - height: image.height as u32, - depth_or_array_layers: 1, - }, - mip_level_count: 1, - sample_count: 1, - dimension: wgpu::TextureDimension::D2, - format: wgpu::TextureFormat::Rgba32Float, - usage: wgpu::TextureUsages::COPY_SRC | wgpu::TextureUsages::STORAGE_BINDING, - view_formats: &[wgpu::TextureFormat::Rgba32Float], - }); + let output_texture = gpu.create_texture(&wgpu::TextureDescriptor { + label: None, + size: wgpu::Extent3d { + width: image.width as u32, + height: image.height as u32, + depth_or_array_layers: 1, + }, + mip_level_count: 1, + sample_count: 1, + dimension: wgpu::TextureDimension::D2, + format: wgpu::TextureFormat::Rgba32Float, + usage: wgpu::TextureUsages::COPY_SRC | wgpu::TextureUsages::STORAGE_BINDING, + view_formats: &[wgpu::TextureFormat::Rgba32Float], + }); - let bind_group = gpu.create_bind_group(&wgpu::BindGroupDescriptor { - label: None, - layout: &self.interpolate_bind_layout, - entries: &[ - wgpu::BindGroupEntry { - binding: 0, - resource: wgpu::BindingResource::TextureView(&input_texture.create_view(&wgpu::TextureViewDescriptor { - label: None, - format: Some(wgpu::TextureFormat::R16Uint), - dimension: Some(wgpu::TextureViewDimension::D2), - usage: Some(wgpu::TextureUsages::STORAGE_BINDING), - aspect: wgpu::TextureAspect::All, - base_mip_level: 0, - mip_level_count: None, - base_array_layer: 0, - array_layer_count: None, - })), - }, - wgpu::BindGroupEntry { - binding: 1, - resource: wgpu::BindingResource::TextureView(&output_texture.create_view(&wgpu::TextureViewDescriptor { - label: None, - format: Some(wgpu::TextureFormat::Rgba32Float), - dimension: Some(wgpu::TextureViewDimension::D2), - usage: Some(wgpu::TextureUsages::STORAGE_BINDING), - aspect: wgpu::TextureAspect::All, - base_mip_level: 0, - mip_level_count: None, - base_array_layer: 0, - array_layer_count: None, - })), - }, - ], - }); + let bind_group = gpu.create_bind_group(&wgpu::BindGroupDescriptor { + label: None, + layout: &self.interpolate_bind_layout, + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: wgpu::BindingResource::TextureView(&input_texture.create_view( + &wgpu::TextureViewDescriptor { + label: None, + format: Some(wgpu::TextureFormat::R16Uint), + dimension: Some(wgpu::TextureViewDimension::D2), + usage: Some(wgpu::TextureUsages::STORAGE_BINDING), + aspect: wgpu::TextureAspect::All, + base_mip_level: 0, + mip_level_count: None, + base_array_layer: 0, + array_layer_count: None, + }, + )), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: wgpu::BindingResource::TextureView(&output_texture.create_view( + &wgpu::TextureViewDescriptor { + label: None, + format: Some(wgpu::TextureFormat::Rgba32Float), + dimension: Some(wgpu::TextureViewDimension::D2), + usage: Some(wgpu::TextureUsages::STORAGE_BINDING), + aspect: wgpu::TextureAspect::All, + base_mip_level: 0, + mip_level_count: None, + base_array_layer: 0, + array_layer_count: None, + }, + )), + }, + ], + }); - let mut encoder = gpu.create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: None, - }); - { - let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { label: None, timestamp_writes: None }); - pass.set_bind_group(0, &bind_group, &[]); - pass.set_pipeline(&self.interpolate_pipeline); - pass.dispatch_workgroups((image.width as u32).div_ceil(8), (image.height as u32).div_ceil(8), 1); - } - let command_buf = encoder.finish(); - queue.submit([command_buf]); + let mut encoder = + gpu.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }); + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: None, + timestamp_writes: None, + }); + pass.set_bind_group(0, &bind_group, &[]); + pass.set_pipeline(&self.interpolate_pipeline); + pass.dispatch_workgroups( + (image.width as u32).div_ceil(8), + (image.height as u32).div_ceil(8), + 1, + ); + } + let command_buf = encoder.finish(); + queue.submit([command_buf]); - output_texture - } + output_texture + } } |