summaryrefslogtreecommitdiff
path: root/src/demosaic
diff options
context:
space:
mode:
Diffstat (limited to 'src/demosaic')
-rw-r--r--src/demosaic/lmmse.rs165
-rw-r--r--src/demosaic/lmmse_01_interpolate.wgsl11
-rw-r--r--src/demosaic/mod.rs7
3 files changed, 183 insertions, 0 deletions
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<r16uint, read>;
+
+@group(0) @binding(1)
+var output_tex: texture_storage_2d<rgba32float, write>;
+
+@compute @workgroup_size(8, 8)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+ let texel = textureLoad(input_tex, global_id.xy);
+ textureStore(output_tex, global_id.xy, vec4<f32>(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;