summary refs log tree commit diff
path: root/src/demosaic/lmmse.rs
blob: 9add4c1dc38d295f99fbc97d3bb9d6f400243f1a (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
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
    }
}