Skip to main content

kuwahara/
kuwahara.rs

1use cuneus::compute::*;
2use cuneus::prelude::*;
3use log::error;
4
5cuneus::uniform_params! {
6    struct KuwaharaParams {
7    radius: f32,
8    q: f32,
9    alpha: f32,
10    filter_strength: f32,
11
12    sigma_d: f32,
13    sigma_r: f32,
14
15    edge_threshold: f32,
16    color_enhance: f32,
17
18    blur_samples: f32,
19    blur_lod: f32,
20    blur_slod: f32,
21
22    filter_mode: i32,
23    show_tensors: i32,
24
25    lic_length: f32,
26    lic_strength: f32,
27    lic_width: f32}
28}
29
30struct KuwaharaShader {
31    base: RenderKit,
32    compute_shader: ComputeShader,
33    current_params: KuwaharaParams}
34
35impl ShaderManager for KuwaharaShader {
36    fn init(core: &Core) -> Self {
37        let initial_params = KuwaharaParams {
38            radius: 5.0,
39            q: 1.5,
40            alpha: 4.0,
41            filter_strength: 0.8,
42            sigma_d: 0.8,
43            sigma_r: 1.2,
44            edge_threshold: 0.2,
45            color_enhance: 1.0,
46            blur_samples: 15.0,
47            blur_lod: 2.0,
48            blur_slod: 4.0,
49            filter_mode: 1,
50            show_tensors: 0,
51            lic_length: 15.0,
52            lic_strength: 0.5,
53            lic_width: 1.5};
54        let base = RenderKit::new(core);
55
56        let passes = vec![
57            PassDescription::new("structure_tensor", &[]),
58            PassDescription::new("tensor_field", &["structure_tensor"]),
59            PassDescription::new("kuwahara_filter", &["tensor_field"]),
60            PassDescription::new("lic_edges", &["tensor_field", "kuwahara_filter"])
61                .with_resolution_scale(0.5),
62            PassDescription::new("main_image", &["lic_edges"]),
63        ];
64
65        let config = ComputeShader::builder()
66            .with_entry_point("structure_tensor")
67            .with_multi_pass(&passes)
68            .with_custom_uniforms::<KuwaharaParams>()
69            .with_workgroup_size([16, 16, 1])
70            .with_texture_format(COMPUTE_TEXTURE_FORMAT_RGBA16)
71            .with_channels(2)
72            .with_label("Kuwahara Multi-Pass")
73            .build();
74
75        let compute_shader = cuneus::compute_shader!(core, "shaders/kuwahara.wgsl", config);
76
77        compute_shader.set_custom_params(initial_params, &core.queue);
78
79        Self {
80            base,
81            compute_shader,
82            current_params: initial_params}
83    }
84
85    fn update(&mut self, core: &Core) {
86        let current_time = self.base.controls.get_time(&self.base.start_time);
87        let delta = 1.0 / 60.0;
88        self.compute_shader
89            .set_time(current_time, delta, &core.queue);
90
91        self.base.update_current_texture(core, &core.queue);
92        if let Some(texture_manager) = self.base.get_current_texture_manager() {
93            self.compute_shader.update_channel_texture(
94                0,
95                &texture_manager.view,
96                &texture_manager.sampler,
97                &core.device,
98                &core.queue,
99            );
100        }
101
102        self.compute_shader.handle_export(core, &mut self.base);
103    }
104
105    fn resize(&mut self, core: &Core) {
106        self.base.default_resize(core, &mut self.compute_shader);
107    }
108
109    fn render(&mut self, core: &Core) -> Result<(), cuneus::SurfaceError> {
110        let mut frame = self.base.begin_frame(core)?;
111
112        let mut params = self.current_params;
113        let mut changed = false;
114        let mut should_start_export = false;
115        let mut export_request = self.base.export_manager.get_ui_request();
116        let mut controls_request = self
117            .base
118            .controls
119            .get_ui_request(&self.base.start_time, &core.size, self.base.fps_tracker.fps());
120
121        let using_video_texture = self.base.using_video_texture;
122        let using_hdri_texture = self.base.using_hdri_texture;
123        let using_webcam_texture = self.base.using_webcam_texture;
124        let video_info = self.base.get_video_info();
125        let hdri_info = self.base.get_hdri_info();
126        let webcam_info = self.base.get_webcam_info();
127
128        let current_fps = self.base.fps_tracker.fps();
129
130        let full_output = if self.base.key_handler.show_ui {
131            self.base.render_ui(core, |ctx| {
132                RenderKit::apply_default_style(ctx);
133
134                egui::Window::new("Filter")
135                    .collapsible(true)
136                    .resizable(true)
137                    .default_width(320.0)
138                    .show(ctx, |ui| {
139                        ShaderControls::render_media_panel(
140                            ui,
141                            &mut controls_request,
142                            using_video_texture,
143                            video_info,
144                            using_hdri_texture,
145                            hdri_info,
146                            using_webcam_texture,
147                            webcam_info,
148                        );
149
150                        ui.separator();
151
152                        let mut anisotropy_enabled = params.filter_mode == 1;
153                        if ui
154                            .checkbox(&mut anisotropy_enabled, "Anisotropy?")
155                            .changed()
156                        {
157                            params.filter_mode = if anisotropy_enabled { 1 } else { 0 };
158                            changed = true;
159                        }
160
161                        egui::CollapsingHeader::new("Filter Parameters")
162                            .default_open(true)
163                            .show(ui, |ui| {
164                                changed |= ui
165                                    .add(
166                                        egui::Slider::new(&mut params.radius, 2.0..=16.0)
167                                            .text("Radius"),
168                                    )
169                                    .changed();
170                                changed |= ui
171                                    .add(
172                                        egui::Slider::new(&mut params.filter_strength, 0.0..=16.0)
173                                            .text("Filter Strength"),
174                                    )
175                                    .changed();
176
177                                if params.filter_mode == 1 {
178                                    ui.separator();
179                                    ui.label("Anisotropic Controls:");
180                                    changed |= ui
181                                        .add(
182                                            egui::Slider::new(&mut params.alpha, 0.1..=16.0)
183                                                .text("Anisotropy"),
184                                        )
185                                        .changed();
186                                }
187                            });
188                        egui::CollapsingHeader::new("Blur Settings")
189                            .default_open(false)
190                            .show(ui, |ui| {
191                                changed |= ui
192                                    .add(
193                                        egui::Slider::new(&mut params.blur_samples, 5.0..=25.0)
194                                            .text("Samples"),
195                                    )
196                                    .changed();
197                                changed |= ui
198                                    .add(
199                                        egui::Slider::new(&mut params.blur_lod, 0.0..=5.0)
200                                            .text("LOD"),
201                                    )
202                                    .changed();
203                                changed |= ui
204                                    .add(
205                                        egui::Slider::new(&mut params.blur_slod, 2.0..=5.0)
206                                            .text("Step"),
207                                    )
208                                    .changed();
209                            });
210
211                        egui::CollapsingHeader::new("Brush Strokes (LIC)")
212                            .default_open(true)
213                            .show(ui, |ui| {
214                                changed |= ui
215                                    .add(
216                                        egui::Slider::new(&mut params.lic_strength, 0.0..=1.0)
217                                            .text("Strength"),
218                                    )
219                                    .changed();
220                                changed |= ui
221                                    .add(
222                                        egui::Slider::new(&mut params.lic_length, 3.0..=40.0)
223                                            .text("Stroke Length"),
224                                    )
225                                    .changed();
226                                changed |= ui
227                                    .add(
228                                        egui::Slider::new(&mut params.lic_width, 0.5..=4.0)
229                                            .text("Stroke Width"),
230                                    )
231                                    .changed();
232                            });
233
234                        egui::CollapsingHeader::new("Post-Processing")
235                            .default_open(false)
236                            .show(ui, |ui| {
237                                changed |= ui
238                                    .add(
239                                        egui::Slider::new(&mut params.color_enhance, 0.5..=2.0)
240                                            .text("Color Filter"),
241                                    )
242                                    .changed();
243
244                                ui.separator();
245                                if ui.button("Reset to Defaults").clicked() {
246                                    params = KuwaharaParams {
247                                        radius: 8.0,
248                                        q: 8.0,
249                                        alpha: 1.0,
250                                        filter_strength: 1.0,
251                                        sigma_d: 1.0,
252                                        sigma_r: 2.0,
253                                        edge_threshold: 0.2,
254                                        color_enhance: 1.0,
255                                        blur_samples: 35.0,
256                                        blur_lod: 2.0,
257                                        blur_slod: 4.0,
258                                        filter_mode: params.filter_mode,
259                                        show_tensors: 0,
260                                        lic_length: 15.0,
261                                        lic_strength: 0.5,
262                                        lic_width: 1.5};
263                                    changed = true;
264                                }
265                            });
266
267                        ui.separator();
268
269                        ShaderControls::render_controls_widget(ui, &mut controls_request);
270
271                        ui.separator();
272
273                        should_start_export =
274                            ExportManager::render_export_ui_widget(ui, &mut export_request);
275
276                        ui.separator();
277                        ui.label(format!(
278                            "Resolution: {}x{}",
279                            core.size.width, core.size.height
280                        ));
281                        ui.label(format!("FPS: {current_fps:.1}"));
282                    });
283            })
284        } else {
285            self.base.render_ui(core, |_ctx| {})
286        };
287
288        self.base.apply_media_requests(core, &controls_request);
289
290        self.base.export_manager.apply_ui_request(export_request);
291        if should_start_export {
292            self.base.export_manager.start_export();
293        }
294
295        if changed {
296            self.current_params = params;
297            self.compute_shader.set_custom_params(params, &core.queue);
298        }
299
300        self.compute_shader.dispatch(&mut frame.encoder, core);
301
302        self.base.renderer.render_to_view(&mut frame.encoder, &frame.view, &self.compute_shader.get_output_texture().bind_group);
303
304        self.base.end_frame(core, frame, full_output);
305
306        Ok(())
307    }
308
309    fn handle_input(&mut self, core: &Core, event: &WindowEvent) -> bool {
310        if self.base.default_handle_input(core, event) {
311            return true;
312        }
313        if let WindowEvent::DroppedFile(path) = event {
314            if let Err(e) = self.base.load_media(core, path) {
315                error!("Failed to load dropped file: {e:?}");
316            }
317            return true;
318        }
319        false
320    }
321}
322
323fn main() -> Result<(), Box<dyn std::error::Error>> {
324    cuneus::gst::init()?;
325    env_logger::init();
326    let (app, event_loop) = ShaderApp::new("Kuwahara Filter", 800, 600);
327
328    app.run(event_loop, KuwaharaShader::init)
329}