Skip to main content

scenecolor/
scenecolor.rs

1use cuneus::compute::*;
2use cuneus::prelude::*;
3use log::error;
4
5cuneus::uniform_params! {
6    struct SceneColorParams {
7    num_segments: f32,
8    palette_height: f32,
9    samples_x: i32,
10    samples_y: i32,
11
12    _pad1: f32,
13    _pad2: f32,
14    _pad3: f32,
15    _pad4: f32}
16}
17
18struct SceneColorShader {
19    base: RenderKit,
20    compute_shader: ComputeShader,
21    current_params: SceneColorParams}
22
23impl SceneColorShader {
24    fn clear_buffers(&mut self, core: &Core) {
25        self.compute_shader.clear_all_buffers(core);
26    }
27}
28
29impl ShaderManager for SceneColorShader {
30    fn init(core: &Core) -> Self {
31        let initial_params = SceneColorParams {
32            num_segments: 16.0,
33            palette_height: 0.2,
34            samples_x: 8,
35            samples_y: 8,
36            _pad1: 0.0,
37            _pad2: 0.0,
38            _pad3: 0.0,
39            _pad4: 0.0};
40
41        let base = RenderKit::new(core);
42
43        let config = ComputeShader::builder()
44            .with_entry_point("main")
45            .with_input_texture() // Enable input texture support
46            .with_custom_uniforms::<SceneColorParams>()
47            .with_workgroup_size([16, 16, 1])
48            .with_texture_format(COMPUTE_TEXTURE_FORMAT_RGBA16)
49            .with_label("Scene Color Unified")
50            .build();
51
52        let compute_shader = cuneus::compute_shader!(core, "shaders/scenecolor.wgsl", config);
53
54        compute_shader.set_custom_params(initial_params, &core.queue);
55
56        Self {
57            base,
58            compute_shader,
59            current_params: initial_params}
60    }
61
62    fn update(&mut self, core: &Core) {
63        // Update time
64        let current_time = self.base.controls.get_time(&self.base.start_time);
65        let delta = 1.0 / 60.0;
66        self.compute_shader
67            .set_time(current_time, delta, &core.queue);
68
69        // Update input textures for media processing
70        self.base.update_current_texture(core, &core.queue);
71        if let Some(texture_manager) = self.base.get_current_texture_manager() {
72            self.compute_shader.update_input_texture(
73                &texture_manager.view,
74                &texture_manager.sampler,
75                &core.device,
76            );
77        }
78    }
79
80    fn resize(&mut self, core: &Core) {
81        self.base.default_resize(core, &mut self.compute_shader);
82    }
83
84    fn render(&mut self, core: &Core) -> Result<(), cuneus::SurfaceError> {
85        let mut frame = self.base.begin_frame(core)?;
86
87        let mut params = self.current_params;
88        let mut changed = false;
89        let mut should_start_export = false;
90        let mut export_request = self.base.export_manager.get_ui_request();
91        let mut controls_request = self
92            .base
93            .controls
94            .get_ui_request(&self.base.start_time, &core.size, self.base.fps_tracker.fps());
95
96        let using_video_texture = self.base.using_video_texture;
97        let using_hdri_texture = self.base.using_hdri_texture;
98        let using_webcam_texture = self.base.using_webcam_texture;
99        let video_info = self.base.get_video_info();
100        let hdri_info = self.base.get_hdri_info();
101        let webcam_info = self.base.get_webcam_info();
102
103        let full_output = if self.base.key_handler.show_ui {
104            self.base.render_ui(core, |ctx| {
105                RenderKit::apply_default_style(ctx);
106
107                egui::Window::new("Scene Color Palette")
108                    .collapsible(true)
109                    .resizable(true)
110                    .default_width(280.0)
111                    .show(ctx, |ui| {
112                        // Media controls
113                        ShaderControls::render_media_panel(
114                            ui,
115                            &mut controls_request,
116                            using_video_texture,
117                            video_info,
118                            using_hdri_texture,
119                            hdri_info,
120                            using_webcam_texture,
121                            webcam_info,
122                        );
123
124                        ui.separator();
125
126                        egui::CollapsingHeader::new("Palette Parameters")
127                            .default_open(true)
128                            .show(ui, |ui| {
129                                changed |= ui
130                                    .add(
131                                        egui::Slider::new(&mut params.num_segments, 1.0..=64.0)
132                                            .text("Segments"),
133                                    )
134                                    .changed();
135                                changed |= ui
136                                    .add(
137                                        egui::Slider::new(&mut params.palette_height, 0.05..=0.5)
138                                            .text("Height"),
139                                    )
140                                    .changed();
141                                changed |= ui
142                                    .add(
143                                        egui::Slider::new(&mut params.samples_x, 1..=32)
144                                            .text("Samples X"),
145                                    )
146                                    .changed();
147                                changed |= ui
148                                    .add(
149                                        egui::Slider::new(&mut params.samples_y, 1..=32)
150                                            .text("Samples Y"),
151                                    )
152                                    .changed();
153                            });
154
155                        ui.separator();
156                        ShaderControls::render_controls_widget(ui, &mut controls_request);
157
158                        ui.separator();
159                        should_start_export =
160                            ExportManager::render_export_ui_widget(ui, &mut export_request);
161
162                        ui.separator();
163                        ui.label("Color palette extractor from scene");
164                    });
165            })
166        } else {
167            self.base.render_ui(core, |_ctx| {})
168        };
169
170        // Apply controls
171        self.base.export_manager.apply_ui_request(export_request);
172        if controls_request.should_clear_buffers {
173            self.clear_buffers(core);
174        }
175        self.base.apply_media_requests(core, &controls_request);
176
177        if changed {
178            self.current_params = params;
179            self.compute_shader.set_custom_params(params, &core.queue);
180        }
181
182        if should_start_export {
183            self.base.export_manager.start_export();
184        }
185
186        // Single stage dispatch
187        self.compute_shader.dispatch(&mut frame.encoder, core);
188
189        self.base.renderer.render_to_view(&mut frame.encoder, &frame.view, &self.compute_shader.get_output_texture().bind_group);
190
191        self.base.end_frame(core, frame, full_output);
192
193        Ok(())
194    }
195
196    fn handle_input(&mut self, core: &Core, event: &WindowEvent) -> bool {
197        if self.base.default_handle_input(core, event) {
198            return true;
199        }
200        if let WindowEvent::DroppedFile(path) = event {
201            if let Err(e) = self.base.load_media(core, path) {
202                error!("Failed to load dropped file: {e:?}");
203            }
204            return true;
205        }
206        false
207    }
208}
209
210fn main() -> Result<(), Box<dyn std::error::Error>> {
211    cuneus::gst::init()?;
212    env_logger::init();
213    let (app, event_loop) = ShaderApp::new("Scene Color Palette", 800, 600);
214    app.run(event_loop, SceneColorShader::init)
215}