From 1d888631c3e4c423c2b782b01ca29fc1057c51a2 Mon Sep 17 00:00:00 2001 From: Maxwell Beck Date: Sun, 30 Mar 2025 15:46:30 -0500 Subject: Initial commit --- src/demosaic/lmmse.rs | 165 +++++++++++++++++++++++++++++++++ src/demosaic/lmmse_01_interpolate.wgsl | 11 +++ src/demosaic/mod.rs | 7 ++ 3 files changed, 183 insertions(+) create mode 100644 src/demosaic/lmmse.rs create mode 100644 src/demosaic/lmmse_01_interpolate.wgsl create mode 100644 src/demosaic/mod.rs (limited to 'src/demosaic') diff --git a/src/demosaic/lmmse.rs b/src/demosaic/lmmse.rs new file mode 100644 index 0000000..1dc3094 --- /dev/null +++ b/src/demosaic/lmmse.rs @@ -0,0 +1,165 @@ +pub struct Lmmse { + 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, + }); + + 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], + }); + + 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 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]); + + output_texture + } +} diff --git a/src/demosaic/lmmse_01_interpolate.wgsl b/src/demosaic/lmmse_01_interpolate.wgsl new file mode 100644 index 0000000..57cd328 --- /dev/null +++ b/src/demosaic/lmmse_01_interpolate.wgsl @@ -0,0 +1,11 @@ +@group(0) @binding(0) +var input_tex: texture_storage_2d; + +@group(0) @binding(1) +var output_tex: texture_storage_2d; + +@compute @workgroup_size(8, 8) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let texel = textureLoad(input_tex, global_id.xy); + textureStore(output_tex, global_id.xy, vec4(texel) / 1024f); +} diff --git a/src/demosaic/mod.rs b/src/demosaic/mod.rs new file mode 100644 index 0000000..4e74f4b --- /dev/null +++ b/src/demosaic/mod.rs @@ -0,0 +1,7 @@ +pub trait Demosaic { + fn new(gpu: &wgpu::Device, queue: &wgpu::Queue) -> Self where Self: Sized; + fn demoasic(&self, gpu: &wgpu::Device, queue: &wgpu::Queue, image: &rawloader::RawImage) -> wgpu::Texture; +} + +pub mod lmmse; +pub use lmmse::Lmmse; -- cgit 1.4.1