From 67feb5c8abfe2ffb2b7ce87552f218ce274dc70c Mon Sep 17 00:00:00 2001
From: Maxwell Beck <max@rastertail.net>
Date: Tue, 8 Apr 2025 18:16:28 -0500
Subject: Refactor `Demosaic` trait

---
 src/demosaic/lmmse.rs | 44 +++++++++++++++++++++++++++++---------------
 src/demosaic/mod.rs   | 14 +++++++-------
 src/main.rs           |  8 +++++---
 3 files changed, 41 insertions(+), 25 deletions(-)

diff --git a/src/demosaic/lmmse.rs b/src/demosaic/lmmse.rs
index 9add4c1..9282380 100644
--- a/src/demosaic/lmmse.rs
+++ b/src/demosaic/lmmse.rs
@@ -1,4 +1,8 @@
 pub struct Lmmse {
+    input_texture: Option<wgpu::Texture>,
+    input_width: u32,
+    input_height: u32,
+
     interpolate_bind_layout: wgpu::BindGroupLayout,
     interpolate_pipeline_layout: wgpu::PipelineLayout,
     interpolate_shader: wgpu::ShaderModule,
@@ -6,7 +10,7 @@ pub struct Lmmse {
 }
 
 impl super::Demosaic for Lmmse {
-    fn new(gpu: &wgpu::Device, queue: &wgpu::Queue) -> Self {
+    fn new(gpu: &wgpu::Device) -> Self {
         let interpolate_bind_layout =
             gpu.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                 label: None,
@@ -53,6 +57,10 @@ impl super::Demosaic for Lmmse {
         });
 
         Self {
+            input_texture: None,
+            input_width: 0,
+            input_height: 0,
+
             interpolate_bind_layout,
             interpolate_pipeline_layout,
             interpolate_shader,
@@ -60,13 +68,10 @@ impl super::Demosaic for Lmmse {
         }
     }
 
-    fn demoasic(
-        &self,
-        gpu: &wgpu::Device,
-        queue: &wgpu::Queue,
-        image: &rawloader::RawImage,
-    ) -> wgpu::Texture {
-        let input_texture = gpu.create_texture(&wgpu::TextureDescriptor {
+    fn bind_image(&mut self, gpu: &wgpu::Device, queue: &wgpu::Queue, image: &rawloader::RawImage) {
+        self.input_width = image.width as u32;
+        self.input_height = image.height as u32;
+        self.input_texture = Some(gpu.create_texture(&wgpu::TextureDescriptor {
             label: None,
             size: wgpu::Extent3d {
                 width: image.width as u32,
@@ -79,7 +84,7 @@ impl super::Demosaic for Lmmse {
             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(),
@@ -87,7 +92,7 @@ impl super::Demosaic for Lmmse {
         };
         queue.write_texture(
             wgpu::TexelCopyTextureInfo {
-                texture: &input_texture,
+                texture: self.input_texture.as_ref().unwrap(),
                 mip_level: 0,
                 origin: wgpu::Origin3d::ZERO,
                 aspect: wgpu::TextureAspect::All,
@@ -104,12 +109,19 @@ impl super::Demosaic for Lmmse {
                 depth_or_array_layers: 1,
             },
         );
+    }
+
+    fn demoasic(&self, gpu: &wgpu::Device, queue: &wgpu::Queue) {
+        let input_texture = self
+            .input_texture
+            .as_ref()
+            .expect("Attempted to demosaic with no input texture bound");
 
         let output_texture = gpu.create_texture(&wgpu::TextureDescriptor {
             label: None,
             size: wgpu::Extent3d {
-                width: image.width as u32,
-                height: image.height as u32,
+                width: self.input_width,
+                height: self.input_height,
                 depth_or_array_layers: 1,
             },
             mip_level_count: 1,
@@ -169,14 +181,16 @@ impl super::Demosaic for Lmmse {
             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),
+                self.input_width.div_ceil(8),
+                self.input_height.div_ceil(8),
                 1,
             );
         }
         let command_buf = encoder.finish();
         queue.submit([command_buf]);
+    }
 
-        output_texture
+    fn get_output(&self) -> &wgpu::Texture {
+        unimplemented!();
     }
 }
diff --git a/src/demosaic/mod.rs b/src/demosaic/mod.rs
index e3d5209..d21d1b5 100644
--- a/src/demosaic/mod.rs
+++ b/src/demosaic/mod.rs
@@ -1,13 +1,13 @@
 pub trait Demosaic {
-    fn new(gpu: &wgpu::Device, queue: &wgpu::Queue) -> Self
+    fn new(gpu: &wgpu::Device) -> Self
     where
         Self: Sized;
-    fn demoasic(
-        &self,
-        gpu: &wgpu::Device,
-        queue: &wgpu::Queue,
-        image: &rawloader::RawImage,
-    ) -> wgpu::Texture;
+
+    fn bind_image(&mut self, gpu: &wgpu::Device, queue: &wgpu::Queue, image: &rawloader::RawImage);
+
+    fn demoasic(&self, gpu: &wgpu::Device, queue: &wgpu::Queue);
+
+    fn get_output(&self) -> &wgpu::Texture;
 }
 
 pub mod lmmse;
diff --git a/src/main.rs b/src/main.rs
index 5775d70..d6abaf4 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -54,10 +54,12 @@ fn main() {
         let image = rawloader::decode_file(path).unwrap();
         dbg!(&image.cfa);
 
-        let pipeline = pipeline::Resources {
-            demosaic: Box::new(demosaic::Lmmse::new(&gpu, &queue)),
+        let mut pipeline = pipeline::Resources {
+            demosaic: Box::new(demosaic::Lmmse::new(&gpu)),
         };
-        let demosaiced = pipeline.demosaic.demoasic(&gpu, &queue, &image);
+        pipeline.demosaic.bind_image(&gpu, &queue, &image);
+        pipeline.demosaic.demoasic(&gpu, &queue);
+        let demosaiced = pipeline.demosaic.get_output();
 
         let readback_buf = gpu.create_buffer(&wgpu::BufferDescriptor {
             label: None,
-- 
cgit 1.4.1