Skip to main content

threecrate_visualization/
interactive_viewer.rs

1//! Interactive 3D viewer with UI controls
2//! 
3//! This module provides a simplified interactive viewer for 3D data
4
5use std::sync::Arc;
6use winit::{
7    application::ApplicationHandler,
8    event::{WindowEvent, ElementState, MouseButton},
9    event_loop::{EventLoop, ActiveEventLoop},
10    window::{Window, WindowId},
11    keyboard::Key,
12    dpi::PhysicalPosition,
13};
14
15use threecrate_core::{PointCloud, TriangleMesh, Result, Point3f, ColoredPoint3f, Error};
16use threecrate_gpu::{
17    PointCloudRenderer, RenderConfig, PointVertex,
18    MeshRenderer, MeshRenderConfig, ShadingMode, PbrMaterial, MeshLightingParams, mesh_to_gpu_mesh,
19};
20use threecrate_algorithms::{ICPResult, PlaneSegmentationResult};
21use crate::camera::Camera;
22
23use nalgebra::{Vector3, Point3};
24
25/// Types of data that can be displayed
26#[derive(Debug, Clone)]
27pub enum ViewData {
28    Empty,
29    PointCloud(PointCloud<Point3f>),
30    ColoredPointCloud(PointCloud<ColoredPoint3f>),
31    Mesh(TriangleMesh),
32}
33
34/// Camera control modes
35#[derive(Debug, Clone, Copy, PartialEq)]
36pub enum CameraMode {
37    Orbit,
38    Pan,
39    Zoom,
40}
41
42/// Pipeline processing type
43#[derive(Debug, Clone, Copy, PartialEq)]
44pub enum PipelineType {
45    Cpu,
46    Gpu,
47}
48
49/// ICP algorithm parameters
50#[derive(Debug, Clone)]
51pub struct ICPParams {
52    pub max_iterations: usize,
53    pub convergence_threshold: f32,
54    pub max_correspondence_distance: f32,
55}
56
57impl Default for ICPParams {
58    fn default() -> Self {
59        Self {
60            max_iterations: 50,
61            convergence_threshold: 0.001,
62            max_correspondence_distance: 1.0,
63        }
64    }
65}
66
67/// RANSAC algorithm parameters
68#[derive(Debug, Clone)]
69pub struct RANSACParams {
70    pub max_iterations: usize,
71    pub distance_threshold: f32,
72}
73
74impl Default for RANSACParams {
75    fn default() -> Self {
76        Self {
77            max_iterations: 1000,
78            distance_threshold: 0.1,
79        }
80    }
81}
82
83/// UI state for all panels and controls (kept for future use)
84#[derive(Debug)]
85pub struct UIState {
86    pub render_panel_open: bool,
87    pub algorithm_panel_open: bool,
88    pub camera_panel_open: bool,
89    pub stats_panel_open: bool,
90    pub icp_params: ICPParams,
91    pub ransac_params: RANSACParams,
92    pub source_cloud: Option<PointCloud<Point3f>>,
93    pub target_cloud: Option<PointCloud<Point3f>>,
94    pub icp_result: Option<ICPResult>,
95    pub ransac_result: Option<PlaneSegmentationResult>,
96}
97
98impl Default for UIState {
99    fn default() -> Self {
100        Self {
101            render_panel_open: false,
102            algorithm_panel_open: false,
103            camera_panel_open: false,
104            stats_panel_open: false,
105            icp_params: ICPParams::default(),
106            ransac_params: RANSACParams::default(),
107            source_cloud: None,
108            target_cloud: None,
109            icp_result: None,
110            ransac_result: None,
111        }
112    }
113}
114
115/// Interactive 3D viewer with comprehensive UI controls
116pub struct InteractiveViewer {
117    current_data: ViewData,
118    camera: Camera,
119    camera_mode: CameraMode,
120    last_mouse_pos: Option<PhysicalPosition<f64>>,
121    mouse_pressed: bool,
122    right_mouse_pressed: bool,
123    debug_frame_count: usize,
124    vertices_dirty: bool,
125    /// Active shading mode for mesh rendering
126    pub shading_mode: ShadingMode,
127    /// Per-mesh PBR material properties
128    pub material: PbrMaterial,
129    /// Lighting parameters (ambient, light position/intensity/color, tone-mapping)
130    pub lighting_params: MeshLightingParams,
131    /// Set to true when lighting_params has changed and must be uploaded to the GPU
132    lighting_dirty: bool,
133}
134
135impl InteractiveViewer {
136    /// Create a new interactive viewer
137    pub fn new() -> Result<Self> {
138        let camera = Camera::new(
139            Point3::new(5.0, 5.0, 5.0),
140            Point3::new(0.0, 0.0, 0.0),
141            Vector3::new(0.0, 1.0, 0.0),
142            45.0,
143            1.0,
144            0.1,
145            100.0,
146        );
147
148        Ok(Self {
149            current_data: ViewData::Empty,
150            camera,
151            camera_mode: CameraMode::Orbit,
152            last_mouse_pos: None,
153            mouse_pressed: false,
154            right_mouse_pressed: false,
155            debug_frame_count: 0,
156            vertices_dirty: true,
157            shading_mode: ShadingMode::Flat,
158            material: PbrMaterial::default(),
159            lighting_params: MeshLightingParams::default(),
160            lighting_dirty: false,
161        })
162    }
163
164    /// Set point cloud data
165    pub fn set_point_cloud(&mut self, cloud: &PointCloud<Point3f>) {
166        self.current_data = ViewData::PointCloud(cloud.clone());
167        self.vertices_dirty = true;
168        println!("Set point cloud with {} points", cloud.len());
169    }
170
171    /// Set colored point cloud data
172    pub fn set_colored_point_cloud(&mut self, cloud: &PointCloud<ColoredPoint3f>) {
173        self.current_data = ViewData::ColoredPointCloud(cloud.clone());
174        self.vertices_dirty = true;
175        println!("Set colored point cloud with {} points", cloud.len());
176    }
177
178    /// Set mesh data
179    pub fn set_mesh(&mut self, mesh: &TriangleMesh) {
180        self.current_data = ViewData::Mesh(mesh.clone());
181        self.vertices_dirty = true;
182        println!("Set mesh with {} vertices and {} faces", mesh.vertices.len(), mesh.faces.len());
183    }
184
185    /// Set the shading mode used when rendering meshes.
186    pub fn set_shading_mode(&mut self, mode: ShadingMode) {
187        self.shading_mode = mode;
188        println!("Shading mode: {:?}", mode);
189    }
190
191    /// Set the PBR material applied to the rendered mesh.
192    pub fn set_material(&mut self, material: PbrMaterial) {
193        self.material = material;
194    }
195
196    /// Set full lighting parameters (ambient, light position/intensity/colour, gamma, exposure).
197    pub fn set_lighting_params(&mut self, params: MeshLightingParams) {
198        self.lighting_params = params;
199        self.lighting_dirty = true;
200    }
201
202    /// Run the interactive viewer
203    pub fn run(self) -> Result<()> {
204        println!("Starting threecrate Interactive Viewer...");
205
206        // Create event loop
207        let event_loop = EventLoop::new().map_err(|e| Error::Io(std::io::Error::new(std::io::ErrorKind::Other, format!("Failed to create event loop: {}", e))))?;
208
209        // Create application handler
210        let mut app = ViewerApp {
211            viewer: self,
212            window: None,
213            point_renderer: None,
214            mesh_renderer: None,
215            cached_vertices: Vec::new(),
216            screenshot_pending: false,
217        };
218
219        // Run the event loop
220        event_loop.run_app(&mut app).map_err(|e| Error::Io(std::io::Error::new(std::io::ErrorKind::Other, format!("Event loop error: {}", e))))?;
221
222        Ok(())
223    }
224}
225
226impl Default for InteractiveViewer {
227    fn default() -> Self {
228        Self::new().expect("Failed to create InteractiveViewer")
229    }
230}
231
232/// Application handler for the viewer
233struct ViewerApp {
234    viewer: InteractiveViewer,
235    window: Option<Arc<Window>>,
236    point_renderer: Option<PointCloudRenderer<'static>>,
237    mesh_renderer: Option<MeshRenderer<'static>>,
238    cached_vertices: Vec<PointVertex>,
239    screenshot_pending: bool,
240}
241
242impl ApplicationHandler for ViewerApp {
243    fn resumed(&mut self, event_loop: &ActiveEventLoop) {
244        if self.window.is_none() {
245            println!("Creating window and initializing renderers...");
246
247            // Create window with attributes
248            let window_attrs = Window::default_attributes()
249                .with_title("threecrate Interactive Viewer")
250                .with_inner_size(winit::dpi::LogicalSize::new(1200.0, 800.0));
251
252            let window = match event_loop.create_window(window_attrs) {
253                Ok(w) => Arc::new(w),
254                Err(e) => {
255                    eprintln!("Failed to create window: {}", e);
256                    event_loop.exit();
257                    return;
258                }
259            };
260
261            // Update camera aspect ratio
262            let size = window.inner_size();
263            self.viewer.camera.aspect_ratio = size.width as f32 / size.height as f32;
264
265            // Leak the Arc to get a 'static reference
266            // This is safe because the window will live for the duration of the program
267            let window_ref: &'static Window = unsafe {
268                std::mem::transmute::<&Window, &'static Window>(window.as_ref())
269            };
270
271            // Initialize renderers using the static reference
272            let pc_config = RenderConfig::default();
273            let point_renderer = match pollster::block_on(PointCloudRenderer::new(window_ref, pc_config)) {
274                Ok(r) => r,
275                Err(e) => {
276                    eprintln!("Failed to create point cloud renderer: {}", e);
277                    event_loop.exit();
278                    return;
279                }
280            };
281
282            let mesh_config = MeshRenderConfig::default();
283            let mesh_renderer = match pollster::block_on(MeshRenderer::new(window_ref, mesh_config)) {
284                Ok(r) => r,
285                Err(e) => {
286                    eprintln!("Failed to create mesh renderer: {}", e);
287                    event_loop.exit();
288                    return;
289                }
290            };
291
292            self.window = Some(window);
293            self.point_renderer = Some(point_renderer);
294            self.mesh_renderer = Some(mesh_renderer);
295
296            println!("Viewer initialized successfully. Window should now be visible.");
297            println!();
298            println!("Controls:");
299            println!("  O / P / Z  — Orbit / Pan / Zoom camera mode");
300            println!("  R          — Reset camera");
301            println!("  M          — Toggle Flat / PBR shading");
302            println!("  S          — Save screenshot (screenshot_<timestamp>.png)");
303            println!("  [ / ]      — Decrease / Increase ambient light strength");
304            println!("  - / =      — Decrease / Increase light intensity");
305            println!();
306        }
307    }
308
309    fn window_event(&mut self, event_loop: &ActiveEventLoop, _id: WindowId, event: WindowEvent) {
310        let Some(window) = &self.window else { return; };
311        let Some(point_renderer) = &mut self.point_renderer else { return; };
312        let Some(mesh_renderer) = &mut self.mesh_renderer else { return; };
313
314        match event {
315            WindowEvent::CloseRequested => {
316                event_loop.exit();
317            }
318            WindowEvent::Resized(new_size) => {
319                point_renderer.resize(new_size);
320                mesh_renderer.resize(new_size);
321                self.viewer.camera.aspect_ratio = new_size.width as f32 / new_size.height as f32;
322            }
323            WindowEvent::MouseInput { state, button, .. } => {
324                match button {
325                    MouseButton::Left => {
326                        self.viewer.mouse_pressed = state == ElementState::Pressed;
327                    }
328                    MouseButton::Right => {
329                        self.viewer.right_mouse_pressed = state == ElementState::Pressed;
330                    }
331                    _ => {}
332                }
333            }
334            WindowEvent::CursorMoved { position, .. } => {
335                if let Some(last_pos) = self.viewer.last_mouse_pos {
336                    let delta_x = position.x - last_pos.x;
337                    let delta_y = position.y - last_pos.y;
338
339                    if self.viewer.mouse_pressed {
340                        match self.viewer.camera_mode {
341                            CameraMode::Orbit => {
342                                self.viewer.camera.orbit(delta_x as f32 * 0.01, delta_y as f32 * 0.01);
343                            }
344                            CameraMode::Pan => {
345                                self.viewer.camera.pan(delta_x as f32 * 0.01, delta_y as f32 * 0.01);
346                            }
347                            _ => {}
348                        }
349                    }
350                }
351                self.viewer.last_mouse_pos = Some(position);
352            }
353            WindowEvent::MouseWheel { delta, .. } => {
354                let scroll_delta = match delta {
355                    winit::event::MouseScrollDelta::LineDelta(_, y) => y,
356                    winit::event::MouseScrollDelta::PixelDelta(pos) => pos.y as f32 / 100.0,
357                };
358                self.viewer.camera.zoom(scroll_delta * 0.1);
359            }
360            WindowEvent::KeyboardInput { event, .. } => {
361                if event.state == ElementState::Pressed {
362                    match &event.logical_key {
363                        Key::Character(c) => {
364                            match c.as_str() {
365                                "o" | "O" => {
366                                    self.viewer.camera_mode = CameraMode::Orbit;
367                                    println!("Switched to Orbit mode");
368                                }
369                                "p" | "P" => {
370                                    self.viewer.camera_mode = CameraMode::Pan;
371                                    println!("Switched to Pan mode");
372                                }
373                                "z" | "Z" => {
374                                    self.viewer.camera_mode = CameraMode::Zoom;
375                                    println!("Switched to Zoom mode");
376                                }
377                                "r" | "R" => {
378                                    self.viewer.camera.reset();
379                                    println!("Reset camera");
380                                }
381                                // Toggle between Flat and PBR shading (M key)
382                                "m" | "M" => {
383                                    self.viewer.shading_mode = match self.viewer.shading_mode {
384                                        ShadingMode::Flat => ShadingMode::Pbr,
385                                        ShadingMode::Pbr => ShadingMode::Flat,
386                                    };
387                                    println!("Shading mode: {:?}", self.viewer.shading_mode);
388                                }
389                                // Screenshot (S key)
390                                "s" | "S" => {
391                                    self.screenshot_pending = true;
392                                    println!("Screenshot requested…");
393                                }
394                                // Ambient strength  [ decrease  ] increase
395                                "[" => {
396                                    self.viewer.lighting_params.ambient_strength =
397                                        (self.viewer.lighting_params.ambient_strength - 0.01).max(0.0);
398                                    self.viewer.lighting_dirty = true;
399                                    println!("Ambient strength: {:.3}", self.viewer.lighting_params.ambient_strength);
400                                }
401                                "]" => {
402                                    self.viewer.lighting_params.ambient_strength =
403                                        (self.viewer.lighting_params.ambient_strength + 0.01).min(1.0);
404                                    self.viewer.lighting_dirty = true;
405                                    println!("Ambient strength: {:.3}", self.viewer.lighting_params.ambient_strength);
406                                }
407                                // Light intensity  - decrease  = increase
408                                "-" => {
409                                    self.viewer.lighting_params.light_intensity =
410                                        (self.viewer.lighting_params.light_intensity - 0.1).max(0.0);
411                                    self.viewer.lighting_dirty = true;
412                                    println!("Light intensity: {:.2}", self.viewer.lighting_params.light_intensity);
413                                }
414                                "=" => {
415                                    self.viewer.lighting_params.light_intensity =
416                                        (self.viewer.lighting_params.light_intensity + 0.1).min(10.0);
417                                    self.viewer.lighting_dirty = true;
418                                    println!("Light intensity: {:.2}", self.viewer.lighting_params.light_intensity);
419                                }
420                                _ => {}
421                            }
422                        }
423                        _ => {}
424                    }
425                }
426            }
427            WindowEvent::RedrawRequested => {
428                // Sync lighting params to GPU when changed
429                if self.viewer.lighting_dirty {
430                    mesh_renderer.update_lighting(self.viewer.lighting_params);
431                    self.viewer.lighting_dirty = false;
432                }
433
434                // Update camera matrices
435                let view_matrix = self.viewer.camera.view_matrix();
436                let proj_matrix = self.viewer.camera.projection_matrix();
437                let camera_pos = self.viewer.camera.position.coords;
438                point_renderer.update_camera(view_matrix, proj_matrix, camera_pos);
439                mesh_renderer.update_camera(view_matrix, proj_matrix, camera_pos);
440
441                // Rebuild quad vertices only when data has changed
442                if self.viewer.vertices_dirty {
443                    self.cached_vertices = match &self.viewer.current_data {
444                        ViewData::PointCloud(cloud) => {
445                            let mut vertices = Vec::with_capacity(cloud.len() * 6);
446                            for point in cloud.iter() {
447                                let size = 0.02;
448                                let pos = [point.x, point.y, point.z];
449                                let color = [1.0, 1.0, 1.0];
450                                let normal = [0.0, 0.0, 1.0];
451
452                                let v1 = PointVertex::from_point(&Point3f::new(pos[0] - size, pos[1] - size, pos[2]), color, 16.0, normal);
453                                let v2 = PointVertex::from_point(&Point3f::new(pos[0] + size, pos[1] - size, pos[2]), color, 16.0, normal);
454                                let v3 = PointVertex::from_point(&Point3f::new(pos[0] + size, pos[1] + size, pos[2]), color, 16.0, normal);
455                                let v4 = PointVertex::from_point(&Point3f::new(pos[0] - size, pos[1] + size, pos[2]), color, 16.0, normal);
456
457                                vertices.push(v1);
458                                vertices.push(v2);
459                                vertices.push(v3);
460                                vertices.push(v1);
461                                vertices.push(v3);
462                                vertices.push(v4);
463                            }
464                            vertices
465                        }
466                        ViewData::ColoredPointCloud(cloud) => {
467                            let mut vertices = Vec::with_capacity(cloud.len() * 6);
468                            for point in cloud.iter() {
469                                let size = 0.02;
470                                let pos = [point.position.x, point.position.y, point.position.z];
471                                let color = [
472                                    point.color[0] as f32 / 255.0,
473                                    point.color[1] as f32 / 255.0,
474                                    point.color[2] as f32 / 255.0,
475                                ];
476                                let normal = [0.0, 0.0, 1.0];
477
478                                let v1 = PointVertex::from_point(&Point3f::new(pos[0] - size, pos[1] - size, pos[2]), color, 16.0, normal);
479                                let v2 = PointVertex::from_point(&Point3f::new(pos[0] + size, pos[1] - size, pos[2]), color, 16.0, normal);
480                                let v3 = PointVertex::from_point(&Point3f::new(pos[0] + size, pos[1] + size, pos[2]), color, 16.0, normal);
481                                let v4 = PointVertex::from_point(&Point3f::new(pos[0] - size, pos[1] + size, pos[2]), color, 16.0, normal);
482
483                                vertices.push(v1);
484                                vertices.push(v2);
485                                vertices.push(v3);
486                                vertices.push(v1);
487                                vertices.push(v3);
488                                vertices.push(v4);
489                            }
490                            vertices
491                        }
492                        ViewData::Mesh(_) | ViewData::Empty => vec![],
493                    };
494                    self.viewer.vertices_dirty = false;
495                }
496
497                // Debug: Print vertex count periodically
498                if !self.cached_vertices.is_empty() {
499                    if self.viewer.debug_frame_count % 60 == 0 {
500                        println!("Rendering {} vertices", self.cached_vertices.len());
501                    }
502                }
503                self.viewer.debug_frame_count += 1;
504
505                match &self.viewer.current_data {
506                    ViewData::PointCloud(_) | ViewData::ColoredPointCloud(_) => {
507                        if self.screenshot_pending {
508                            self.screenshot_pending = false;
509                            println!("Screenshot not yet supported for point clouds.");
510                        }
511                        if !self.cached_vertices.is_empty() {
512                            if let Err(e) = point_renderer.render(&self.cached_vertices) {
513                                eprintln!("Render error: {}", e);
514                            }
515                        }
516                    }
517                    ViewData::Mesh(mesh) => {
518                        if !mesh.vertices.is_empty() && !mesh.faces.is_empty() {
519                            let indices: Vec<u32> = mesh
520                                .faces
521                                .iter()
522                                .flat_map(|f| [f[0] as u32, f[1] as u32, f[2] as u32])
523                                .collect();
524
525                            let normals_opt = mesh.normals.as_ref().map(|n| n.as_slice());
526
527                            let colors_f32: Option<Vec<[f32; 3]>> = mesh.colors.as_ref().map(|cols| {
528                                cols.iter()
529                                    .map(|c| [c[0] as f32 / 255.0, c[1] as f32 / 255.0, c[2] as f32 / 255.0])
530                                    .collect()
531                            });
532                            let colors_opt = colors_f32.as_ref().map(|c| c.as_slice());
533
534                            // Use the viewer's material and shading mode
535                            let gpu_mesh = mesh_to_gpu_mesh(
536                                &mesh.vertices,
537                                &indices,
538                                normals_opt,
539                                colors_opt,
540                                Some(self.viewer.material),
541                            );
542
543                            if let Err(e) = mesh_renderer.render(&gpu_mesh, self.viewer.shading_mode) {
544                                eprintln!("Mesh render error: {}", e);
545                            }
546
547                            // Handle screenshot request
548                            if self.screenshot_pending {
549                                self.screenshot_pending = false;
550                                match mesh_renderer.render_to_texture(&gpu_mesh, self.viewer.shading_mode) {
551                                    Ok((pixels, format, w, h)) => {
552                                        save_screenshot(&pixels, format, w, h);
553                                    }
554                                    Err(e) => eprintln!("Screenshot render error: {}", e),
555                                }
556                            }
557                        }
558                    }
559                    ViewData::Empty => {}
560                }
561
562                // Request next frame
563                window.request_redraw();
564            }
565            _ => {}
566        }
567    }
568
569    fn about_to_wait(&mut self, _event_loop: &ActiveEventLoop) {
570        if let Some(window) = &self.window {
571            window.request_redraw();
572        }
573    }
574}
575
576/// Encode raw GPU pixel bytes as a PNG file.
577///
578/// Handles both RGBA and BGRA surface formats transparently.
579fn save_screenshot(pixels: &[u8], format: wgpu::TextureFormat, width: u32, height: u32) {
580    use wgpu::TextureFormat;
581
582    // Convert BGRA → RGBA if the surface uses a BGRA format (common on macOS/Metal)
583    let rgba: Vec<u8> = match format {
584        TextureFormat::Bgra8Unorm | TextureFormat::Bgra8UnormSrgb => {
585            let mut buf = Vec::with_capacity(pixels.len());
586            for chunk in pixels.chunks_exact(4) {
587                buf.push(chunk[2]); // R ← B
588                buf.push(chunk[1]); // G
589                buf.push(chunk[0]); // B ← R
590                buf.push(chunk[3]); // A
591            }
592            buf
593        }
594        _ => pixels.to_vec(),
595    };
596
597    let timestamp = std::time::SystemTime::now()
598        .duration_since(std::time::UNIX_EPOCH)
599        .map(|d| d.as_secs())
600        .unwrap_or(0);
601    let path = format!("screenshot_{}.png", timestamp);
602
603    match image::RgbaImage::from_raw(width, height, rgba) {
604        Some(img) => match img.save(&path) {
605            Ok(()) => println!("Screenshot saved: {}", path),
606            Err(e) => eprintln!("Failed to save screenshot: {}", e),
607        },
608        None => eprintln!("Failed to build image buffer for screenshot"),
609    }
610}
611