diff options
| -rw-r--r-- | flake.nix | 2 | ||||
| -rw-r--r-- | src/demosaic/lmmse.rs | 323 | ||||
| -rw-r--r-- | src/demosaic/mod.rs | 11 | ||||
| -rw-r--r-- | src/main.rs | 51 | ||||
| -rw-r--r-- | src/pipeline.rs | 2 |
5 files changed, 215 insertions, 174 deletions
diff --git a/flake.nix b/flake.nix index f425c9b..e6db1d0 100644 --- a/flake.nix +++ b/flake.nix @@ -9,7 +9,7 @@ outputs = { self, nixpkgs, fenix }: let - system = "aarch64-darwin"; + system = "x86_64-linux"; pkgs = nixpkgs.legacyPackages.${system}; fpkgs = fenix.packages.${system}; rust = with fpkgs; combine (with stable; [ 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 + } } diff --git a/src/demosaic/mod.rs b/src/demosaic/mod.rs index 4e74f4b..e3d5209 100644 --- a/src/demosaic/mod.rs +++ b/src/demosaic/mod.rs @@ -1,6 +1,13 @@ 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; + 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; diff --git a/src/main.rs b/src/main.rs index 53f3a41..5775d70 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,6 @@ use std::path::PathBuf; -use bpaf::{construct, positional, OptionParser, Parser}; +use bpaf::{OptionParser, Parser, construct, positional}; use demosaic::Demosaic; use image::buffer::ConvertBuffer; @@ -9,7 +9,7 @@ mod pipeline; #[derive(Clone, Debug)] struct Args { - paths: Vec<PathBuf> + paths: Vec<PathBuf>, } fn args() -> OptionParser<Args> { @@ -26,19 +26,28 @@ fn main() { // Initialize GPU let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default()); - let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default())).unwrap(); + 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) { + 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(); + 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 { @@ -46,7 +55,7 @@ fn main() { dbg!(&image.cfa); let pipeline = pipeline::Resources { - demosaic: Box::new(demosaic::Lmmse::new(&gpu, &queue)) + demosaic: Box::new(demosaic::Lmmse::new(&gpu, &queue)), }; let demosaiced = pipeline.demosaic.demoasic(&gpu, &queue, &image); @@ -56,9 +65,8 @@ fn main() { usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, mapped_at_creation: false, }); - let mut encoder = gpu.create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: None, - }); + let mut encoder = + gpu.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }); encoder.copy_texture_to_buffer( wgpu::TexelCopyTextureInfo { texture: &demosaiced, @@ -78,7 +86,7 @@ fn main() { width: image.width as u32, height: image.height as u32, depth_or_array_layers: 1, - } + }, ); queue.submit([encoder.finish()]); @@ -87,8 +95,17 @@ fn main() { 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(); + 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 index 3ef02d6..c68935d 100644 --- a/src/pipeline.rs +++ b/src/pipeline.rs @@ -1,5 +1,5 @@ use crate::demosaic::Demosaic; pub struct Resources { - pub demosaic: Box<dyn Demosaic>, + pub demosaic: Box<dyn Demosaic>, } |