Skip to main content

computecolors/
computecolors.rs

1use cuneus::compute::*;
2use cuneus::prelude::*;
3use log::error;
4
5cuneus::uniform_params! {
6    struct SplattingParams {
7    animation_speed: f32,
8    splat_size: f32,
9    particle_spread: f32,
10    intensity: f32,
11    particle_density: f32,
12    brightness: f32,
13    physics_strength: f32,
14    trail_length: f32,
15    trail_decay: f32,
16    flow_strength: f32,
17    _padding1: f32,
18    _padding2: u32}
19}
20
21struct ColorProjection {
22    base: RenderKit,
23    compute_shader: ComputeShader,
24    current_params: SplattingParams}
25
26impl ColorProjection {
27    fn clear_buffers(&mut self, core: &Core) {
28        self.compute_shader.clear_all_buffers(core);
29    }
30}
31
32impl ShaderManager for ColorProjection {
33    fn init(core: &Core) -> Self {
34        let base = RenderKit::new(core);
35
36        let initial_params = SplattingParams {
37            animation_speed: 1.0,
38            splat_size: 0.8,
39            particle_spread: 0.0,
40            intensity: 2.0,
41            particle_density: 0.4,
42            brightness: 36.0,
43            physics_strength: 0.5,
44            trail_length: 0.0,
45            trail_decay: 0.95,
46            flow_strength: 1.0,
47            _padding1: 0.0,
48            _padding2: 0};
49
50        // Define the multi-stage passes
51        let passes = vec![
52            PassDescription::new("clear_buffer", &[]), // Stage 0: Clear atomic buffer
53            PassDescription::new("project_colors", &[]), // Stage 1: Project colors to 3D space
54            PassDescription::new("generate_image", &[]), // Stage 2: Generate final image
55        ];
56
57        let config = ComputeShader::builder()
58            .with_entry_point("clear_buffer")
59            .with_multi_pass(&passes)
60            .with_input_texture() // Enable input texture support
61            .with_custom_uniforms::<SplattingParams>()
62            .with_storage_buffer(StorageBufferSpec::new(
63                "atomic_buffer",
64                (core.size.width * core.size.height * 4 * 4) as u64,
65            ))
66            .with_workgroup_size([16, 16, 1])
67            .with_texture_format(COMPUTE_TEXTURE_FORMAT_RGBA16)
68            .with_label("Particle Splatting Multi-Pass")
69            .build();
70
71        let compute_shader = cuneus::compute_shader!(core, "shaders/computecolors.wgsl", config);
72
73        compute_shader.set_custom_params(initial_params, &core.queue);
74
75        Self {
76            base,
77            compute_shader,
78            current_params: initial_params}
79    }
80
81    fn update(&mut self, core: &Core) {
82        // Update time
83        let current_time = self.base.controls.get_time(&self.base.start_time);
84        let delta = 1.0 / 60.0;
85        self.compute_shader
86            .set_time(current_time, delta, &core.queue);
87
88        // Update input textures for media processing
89        self.base.update_current_texture(core, &core.queue);
90        if let Some(texture_manager) = self.base.get_current_texture_manager() {
91            self.compute_shader.update_input_texture(
92                &texture_manager.view,
93                &texture_manager.sampler,
94                &core.device,
95            );
96        }
97    }
98
99    fn resize(&mut self, core: &Core) {
100        self.base.default_resize(core, &mut self.compute_shader);
101    }
102
103    fn render(&mut self, core: &Core) -> Result<(), cuneus::SurfaceError> {
104        let mut frame = self.base.begin_frame(core)?;
105
106        // Handle UI and controls
107        let mut params = self.current_params;
108        let mut changed = false;
109        let mut should_start_export = false;
110        let mut export_request = self.base.export_manager.get_ui_request();
111        let mut controls_request = self
112            .base
113            .controls
114            .get_ui_request(&self.base.start_time, &core.size, self.base.fps_tracker.fps());
115
116        let using_video_texture = self.base.using_video_texture;
117        let using_hdri_texture = self.base.using_hdri_texture;
118        let using_webcam_texture = self.base.using_webcam_texture;
119        let video_info = self.base.get_video_info();
120        let hdri_info = self.base.get_hdri_info();
121        let webcam_info = self.base.get_webcam_info();
122        let full_output = if self.base.key_handler.show_ui {
123            self.base.render_ui(core, |ctx| {
124                RenderKit::apply_default_style(ctx);
125
126                egui::Window::new("Particle Splatting")
127                    .collapsible(true)
128                    .resizable(true)
129                    .default_width(250.0)
130                    .show(ctx, |ui| {
131                        ShaderControls::render_media_panel(
132                            ui,
133                            &mut controls_request,
134                            using_video_texture,
135                            video_info,
136                            using_hdri_texture,
137                            hdri_info,
138                            using_webcam_texture,
139                            webcam_info,
140                        );
141
142                        ui.separator();
143
144                        egui::CollapsingHeader::new("Particles")
145                            .default_open(true)
146                            .show(ui, |ui| {
147                                changed |= ui
148                                    .add(
149                                        egui::Slider::new(&mut params.particle_density, 0.1..=1.0)
150                                            .text("Density"),
151                                    )
152                                    .changed();
153                                changed |= ui
154                                    .add(
155                                        egui::Slider::new(&mut params.splat_size, 0.1..=2.0)
156                                            .text("Splat Size"),
157                                    )
158                                    .changed();
159                                changed |= ui
160                                    .add(
161                                        egui::Slider::new(&mut params.intensity, 0.1..=6.0)
162                                            .text("Intensity"),
163                                    )
164                                    .changed();
165                                changed |= ui
166                                    .add(
167                                        egui::Slider::new(&mut params.brightness, 36.0..=48.0)
168                                            .text("Brightness"),
169                                    )
170                                    .changed();
171                            });
172
173                        egui::CollapsingHeader::new("Effects")
174                            .default_open(true)
175                            .show(ui, |ui| {
176                                changed |= ui
177                                    .add(
178                                        egui::Slider::new(&mut params.animation_speed, 0.0..=3.0)
179                                            .text("Speed"),
180                                    )
181                                    .changed();
182                                changed |= ui
183                                    .add(
184                                        egui::Slider::new(&mut params.particle_spread, 0.0..=1.0)
185                                            .text("Scramble Amount"),
186                                    )
187                                    .changed();
188                                changed |= ui
189                                    .add(
190                                        egui::Slider::new(&mut params.physics_strength, 0.0..=12.0)
191                                            .text("Return Force"),
192                                    )
193                                    .changed();
194                            });
195
196                        egui::CollapsingHeader::new("Flow Trails")
197                            .default_open(true)
198                            .show(ui, |ui| {
199                                changed |= ui
200                                    .add(
201                                        egui::Slider::new(&mut params.trail_length, 0.0..=2.0)
202                                            .text("Trail Length"),
203                                    )
204                                    .changed();
205                                changed |= ui
206                                    .add(
207                                        egui::Slider::new(&mut params.trail_decay, 0.8..=1.0)
208                                            .text("Trail Decay"),
209                                    )
210                                    .changed();
211                                changed |= ui
212                                    .add(
213                                        egui::Slider::new(&mut params.flow_strength, 0.0..=3.0)
214                                            .text("Flow Strength"),
215                                    )
216                                    .changed();
217                            });
218
219                        ui.separator();
220
221                        ShaderControls::render_controls_widget(ui, &mut controls_request);
222
223                        ui.separator();
224
225                        should_start_export =
226                            ExportManager::render_export_ui_widget(ui, &mut export_request);
227                    });
228            })
229        } else {
230            self.base.render_ui(core, |_ctx| {})
231        };
232
233        // Apply controls
234        self.base.export_manager.apply_ui_request(export_request);
235        if controls_request.should_clear_buffers {
236            self.clear_buffers(core);
237        }
238        self.base.apply_media_requests(core, &controls_request);
239
240        if should_start_export {
241            self.base.export_manager.start_export();
242        }
243
244        if changed {
245            self.current_params = params;
246            self.compute_shader.set_custom_params(params, &core.queue);
247        }
248        // Handle export
249        self.compute_shader.handle_export(core, &mut self.base);
250
251        // Color projection multi-stage dispatch - run all stages every frame for animation
252
253        // Stage 0: Clear atomic buffer (16x16 workgroups)
254        self.compute_shader.dispatch_stage(&mut frame.encoder, core, 0);
255
256        // Stage 1: Project colors to 3D space (uses input texture dimensions)
257        if let Some(texture_manager) = self.base.get_current_texture_manager() {
258            let input_workgroups = [
259                texture_manager.texture.width().div_ceil(16),
260                texture_manager.texture.height().div_ceil(16),
261                1,
262            ];
263            self.compute_shader
264                .dispatch_stage_with_workgroups(&mut frame.encoder, 1, input_workgroups);
265        } else {
266            // Fallback to screen size if no input texture
267            self.compute_shader.dispatch_stage(&mut frame.encoder, core, 1);
268        }
269
270        // Stage 2: Generate final image (16x16 workgroups, screen size)
271        self.compute_shader.dispatch_stage(&mut frame.encoder, core, 2);
272
273        self.base.renderer.render_to_view(&mut frame.encoder, &frame.view, &self.compute_shader.get_output_texture().bind_group);
274
275        self.base.end_frame(core, frame, full_output);
276
277        Ok(())
278    }
279
280    fn handle_input(&mut self, core: &Core, event: &WindowEvent) -> bool {
281        if self.base.default_handle_input(core, event) {
282            return true;
283        }
284        if let WindowEvent::DroppedFile(path) = event {
285            if let Err(e) = self.base.load_media(core, path) {
286                error!("Failed to load dropped file: {e:?}");
287            }
288            return true;
289        }
290        false
291    }
292}
293
294fn main() -> Result<(), Box<dyn std::error::Error>> {
295    cuneus::gst::init()?;
296    env_logger::init();
297    let (app, event_loop) = cuneus::ShaderApp::new("Particle Splatting", 800, 600);
298    app.run(event_loop, ColorProjection::init)
299}