Initial commit

5 files changed, 283 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;
diff --git a/src/main.rs b/src/main.rs
new file mode 100644
index 0000000..53f3a41
--- /dev/null
+++ b/src/main.rs
@@ -0,0 +1,95 @@
+use std::path::PathBuf;
+
+use bpaf::{construct, positional, OptionParser, Parser};
+use demosaic::Demosaic;
+use image::buffer::ConvertBuffer;
+
+mod demosaic;
+mod pipeline;
+
+#[derive(Clone, Debug)]
+struct Args {
+    paths: Vec<PathBuf>
+}
+
+fn args() -> OptionParser<Args> {
+    let paths = positional("FILE").some("must process at least one image");
+
+    construct!(Args { paths })
+        .to_options()
+        .descr("Intuitive raw photo processing engine")
+}
+
+fn main() {
+    // Parse arguments
+    let args = args().fallback_to_usage().run();
+
+    // Initialize GPU
+    let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
+    let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default())).unwrap();
+
+    let downlevel_caps = adapter.get_downlevel_capabilities();
+    if !downlevel_caps.flags.contains(wgpu::DownlevelFlags::COMPUTE_SHADERS) {
+        panic!("GPU does not support compute");
+    }
+
+    let (gpu, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
+        label: None,
+        required_features: wgpu::Features::TEXTURE_ADAPTER_SPECIFIC_FORMAT_FEATURES,
+        required_limits: wgpu::Limits::default(),
+        memory_hints: wgpu::MemoryHints::MemoryUsage,
+    }, None)).unwrap();
+
+    // Process images
+    for path in &args.paths {
+        let image = rawloader::decode_file(path).unwrap();
+        dbg!(&image.cfa);
+
+        let pipeline = pipeline::Resources {
+            demosaic: Box::new(demosaic::Lmmse::new(&gpu, &queue))
+        };
+        let demosaiced = pipeline.demosaic.demoasic(&gpu, &queue, &image);
+
+        let readback_buf = gpu.create_buffer(&wgpu::BufferDescriptor {
+            label: None,
+            size: 4 * 4 * image.width as u64 * image.height as u64,
+            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
+            mapped_at_creation: false,
+        });
+        let mut encoder = gpu.create_command_encoder(&wgpu::CommandEncoderDescriptor {
+            label: None,
+        });
+        encoder.copy_texture_to_buffer(
+            wgpu::TexelCopyTextureInfo {
+                texture: &demosaiced,
+                mip_level: 0,
+                origin: wgpu::Origin3d::ZERO,
+                aspect: wgpu::TextureAspect::All,
+            },
+            wgpu::TexelCopyBufferInfo {
+                buffer: &readback_buf,
+                layout: wgpu::TexelCopyBufferLayout {
+                    offset: 0,
+                    bytes_per_row: Some(4 * 4 * 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,
+            }
+        );
+        queue.submit([encoder.finish()]);
+
+        let readback_slice = readback_buf.slice(..);
+        readback_slice.map_async(wgpu::MapMode::Read, |_| {});
+        gpu.poll(wgpu::Maintain::Wait);
+        {
+            let readback_data = readback_slice.get_mapped_range();
+            let result_image = image::ImageBuffer::<image::Rgba<f32>, _>::from_raw(image.width as u32, image.height as u32, bytemuck::cast_slice(&readback_data)).unwrap();
+            <_ as ConvertBuffer<image::ImageBuffer<image::Rgb<u16>, Vec<u16>>>>::convert(&result_image).save_with_format("out.png", image::ImageFormat::Png).unwrap();
+        }
+        readback_buf.unmap();
+    }
+}
diff --git a/src/pipeline.rs b/src/pipeline.rs
new file mode 100644
index 0000000..3ef02d6
--- /dev/null
+++ b/src/pipeline.rs
@@ -0,0 +1,5 @@
+use crate::demosaic::Demosaic;
+
+pub struct Resources {
+	pub demosaic: Box<dyn Demosaic>,
+}