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    dpi::PhysicalPosition,
9    event::{ElementState, MouseButton, WindowEvent},
10    event_loop::{ActiveEventLoop, EventLoop},
11    keyboard::Key,
12    window::{Window, WindowId},
13};
14
15use crate::camera::Camera;
16use threecrate_algorithms::{ICPResult, PlaneSegmentationResult};
17use threecrate_core::{ColoredPoint3f, Error, Point3f, PointCloud, Result, TriangleMesh};
18use threecrate_gpu::{
19    mesh_to_gpu_mesh, MeshLightingParams, MeshRenderConfig, MeshRenderer, PbrMaterial,
20    PointCloudRenderer, PointVertex, RenderConfig, ShadingMode,
21};
22
23use nalgebra::{Point3, Vector3};
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!(
183            "Set mesh with {} vertices and {} faces",
184            mesh.vertices.len(),
185            mesh.faces.len()
186        );
187    }
188
189    /// Set the shading mode used when rendering meshes.
190    pub fn set_shading_mode(&mut self, mode: ShadingMode) {
191        self.shading_mode = mode;
192        println!("Shading mode: {:?}", mode);
193    }
194
195    /// Set the PBR material applied to the rendered mesh.
196    pub fn set_material(&mut self, material: PbrMaterial) {
197        self.material = material;
198    }
199
200    /// Set full lighting parameters (ambient, light position/intensity/colour, gamma, exposure).
201    pub fn set_lighting_params(&mut self, params: MeshLightingParams) {
202        self.lighting_params = params;
203        self.lighting_dirty = true;
204    }
205
206    /// Run the interactive viewer
207    pub fn run(self) -> Result<()> {
208        println!("Starting threecrate Interactive Viewer...");
209
210        // Create event loop
211        let event_loop = EventLoop::new().map_err(|e| {
212            Error::Io(std::io::Error::new(
213                std::io::ErrorKind::Other,
214                format!("Failed to create event loop: {}", e),
215            ))
216        })?;
217
218        // Create application handler
219        let mut app = ViewerApp {
220            viewer: self,
221            window: None,
222            point_renderer: None,
223            mesh_renderer: None,
224            cached_vertices: Vec::new(),
225            screenshot_pending: false,
226        };
227
228        // Run the event loop
229        event_loop.run_app(&mut app).map_err(|e| {
230            Error::Io(std::io::Error::new(
231                std::io::ErrorKind::Other,
232                format!("Event loop error: {}", e),
233            ))
234        })?;
235
236        Ok(())
237    }
238}
239
240impl Default for InteractiveViewer {
241    fn default() -> Self {
242        Self::new().expect("Failed to create InteractiveViewer")
243    }
244}
245
246/// Application handler for the viewer
247struct ViewerApp {
248    viewer: InteractiveViewer,
249    window: Option<Arc<Window>>,
250    point_renderer: Option<PointCloudRenderer<'static>>,
251    mesh_renderer: Option<MeshRenderer<'static>>,
252    cached_vertices: Vec<PointVertex>,
253    screenshot_pending: bool,
254}
255
256impl ApplicationHandler for ViewerApp {
257    fn resumed(&mut self, event_loop: &ActiveEventLoop) {
258        if self.window.is_none() {
259            println!("Creating window and initializing renderers...");
260
261            // Create window with attributes
262            let window_attrs = Window::default_attributes()
263                .with_title("threecrate Interactive Viewer")
264                .with_inner_size(winit::dpi::LogicalSize::new(1200.0, 800.0));
265
266            let window = match event_loop.create_window(window_attrs) {
267                Ok(w) => Arc::new(w),
268                Err(e) => {
269                    eprintln!("Failed to create window: {}", e);
270                    event_loop.exit();
271                    return;
272                }
273            };
274
275            // Update camera aspect ratio
276            let size = window.inner_size();
277            self.viewer.camera.aspect_ratio = size.width as f32 / size.height as f32;
278
279            // Leak the Arc to get a 'static reference
280            // This is safe because the window will live for the duration of the program
281            let window_ref: &'static Window =
282                unsafe { std::mem::transmute::<&Window, &'static Window>(window.as_ref()) };
283
284            // Initialize renderers using the static reference
285            let pc_config = RenderConfig::default();
286            let point_renderer =
287                match pollster::block_on(PointCloudRenderer::new(window_ref, pc_config)) {
288                    Ok(r) => r,
289                    Err(e) => {
290                        eprintln!("Failed to create point cloud renderer: {}", e);
291                        event_loop.exit();
292                        return;
293                    }
294                };
295
296            let mesh_config = MeshRenderConfig::default();
297            let mesh_renderer = match pollster::block_on(MeshRenderer::new(window_ref, mesh_config))
298            {
299                Ok(r) => r,
300                Err(e) => {
301                    eprintln!("Failed to create mesh renderer: {}", e);
302                    event_loop.exit();
303                    return;
304                }
305            };
306
307            self.window = Some(window);
308            self.point_renderer = Some(point_renderer);
309            self.mesh_renderer = Some(mesh_renderer);
310
311            println!("Viewer initialized successfully. Window should now be visible.");
312            println!();
313            println!("Controls:");
314            println!("  O / P / Z  — Orbit / Pan / Zoom camera mode");
315            println!("  R          — Reset camera");
316            println!("  M          — Toggle Flat / PBR shading");
317            println!("  S          — Save screenshot (screenshot_<timestamp>.png)");
318            println!("  [ / ]      — Decrease / Increase ambient light strength");
319            println!("  - / =      — Decrease / Increase light intensity");
320            println!();
321        }
322    }
323
324    fn window_event(&mut self, event_loop: &ActiveEventLoop, _id: WindowId, event: WindowEvent) {
325        let Some(window) = &self.window else {
326            return;
327        };
328        let Some(point_renderer) = &mut self.point_renderer else {
329            return;
330        };
331        let Some(mesh_renderer) = &mut self.mesh_renderer else {
332            return;
333        };
334
335        match event {
336            WindowEvent::CloseRequested => {
337                event_loop.exit();
338            }
339            WindowEvent::Resized(new_size) => {
340                point_renderer.resize(new_size);
341                mesh_renderer.resize(new_size);
342                self.viewer.camera.aspect_ratio = new_size.width as f32 / new_size.height as f32;
343            }
344            WindowEvent::MouseInput { state, button, .. } => match button {
345                MouseButton::Left => {
346                    self.viewer.mouse_pressed = state == ElementState::Pressed;
347                }
348                MouseButton::Right => {
349                    self.viewer.right_mouse_pressed = state == ElementState::Pressed;
350                }
351                _ => {}
352            },
353            WindowEvent::CursorMoved { position, .. } => {
354                if let Some(last_pos) = self.viewer.last_mouse_pos {
355                    let delta_x = position.x - last_pos.x;
356                    let delta_y = position.y - last_pos.y;
357
358                    if self.viewer.mouse_pressed {
359                        match self.viewer.camera_mode {
360                            CameraMode::Orbit => {
361                                self.viewer
362                                    .camera
363                                    .orbit(delta_x as f32 * 0.01, delta_y as f32 * 0.01);
364                            }
365                            CameraMode::Pan => {
366                                self.viewer
367                                    .camera
368                                    .pan(delta_x as f32 * 0.01, delta_y as f32 * 0.01);
369                            }
370                            _ => {}
371                        }
372                    }
373                }
374                self.viewer.last_mouse_pos = Some(position);
375            }
376            WindowEvent::MouseWheel { delta, .. } => {
377                let scroll_delta = match delta {
378                    winit::event::MouseScrollDelta::LineDelta(_, y) => y,
379                    winit::event::MouseScrollDelta::PixelDelta(pos) => pos.y as f32 / 100.0,
380                };
381                self.viewer.camera.zoom(scroll_delta * 0.1);
382            }
383            WindowEvent::KeyboardInput { event, .. } => {
384                if event.state == ElementState::Pressed {
385                    match &event.logical_key {
386                        Key::Character(c) => {
387                            match c.as_str() {
388                                "o" | "O" => {
389                                    self.viewer.camera_mode = CameraMode::Orbit;
390                                    println!("Switched to Orbit mode");
391                                }
392                                "p" | "P" => {
393                                    self.viewer.camera_mode = CameraMode::Pan;
394                                    println!("Switched to Pan mode");
395                                }
396                                "z" | "Z" => {
397                                    self.viewer.camera_mode = CameraMode::Zoom;
398                                    println!("Switched to Zoom mode");
399                                }
400                                "r" | "R" => {
401                                    self.viewer.camera.reset();
402                                    println!("Reset camera");
403                                }
404                                // Toggle between Flat and PBR shading (M key)
405                                "m" | "M" => {
406                                    self.viewer.shading_mode = match self.viewer.shading_mode {
407                                        ShadingMode::Flat => ShadingMode::Pbr,
408                                        ShadingMode::Pbr => ShadingMode::Flat,
409                                    };
410                                    println!("Shading mode: {:?}", self.viewer.shading_mode);
411                                }
412                                // Screenshot (S key)
413                                "s" | "S" => {
414                                    self.screenshot_pending = true;
415                                    println!("Screenshot requested…");
416                                }
417                                // Ambient strength  [ decrease  ] increase
418                                "[" => {
419                                    self.viewer.lighting_params.ambient_strength =
420                                        (self.viewer.lighting_params.ambient_strength - 0.01)
421                                            .max(0.0);
422                                    self.viewer.lighting_dirty = true;
423                                    println!(
424                                        "Ambient strength: {:.3}",
425                                        self.viewer.lighting_params.ambient_strength
426                                    );
427                                }
428                                "]" => {
429                                    self.viewer.lighting_params.ambient_strength =
430                                        (self.viewer.lighting_params.ambient_strength + 0.01)
431                                            .min(1.0);
432                                    self.viewer.lighting_dirty = true;
433                                    println!(
434                                        "Ambient strength: {:.3}",
435                                        self.viewer.lighting_params.ambient_strength
436                                    );
437                                }
438                                // Light intensity  - decrease  = increase
439                                "-" => {
440                                    self.viewer.lighting_params.light_intensity =
441                                        (self.viewer.lighting_params.light_intensity - 0.1)
442                                            .max(0.0);
443                                    self.viewer.lighting_dirty = true;
444                                    println!(
445                                        "Light intensity: {:.2}",
446                                        self.viewer.lighting_params.light_intensity
447                                    );
448                                }
449                                "=" => {
450                                    self.viewer.lighting_params.light_intensity =
451                                        (self.viewer.lighting_params.light_intensity + 0.1)
452                                            .min(10.0);
453                                    self.viewer.lighting_dirty = true;
454                                    println!(
455                                        "Light intensity: {:.2}",
456                                        self.viewer.lighting_params.light_intensity
457                                    );
458                                }
459                                _ => {}
460                            }
461                        }
462                        _ => {}
463                    }
464                }
465            }
466            WindowEvent::RedrawRequested => {
467                // Sync lighting params to GPU when changed
468                if self.viewer.lighting_dirty {
469                    mesh_renderer.update_lighting(self.viewer.lighting_params);
470                    self.viewer.lighting_dirty = false;
471                }
472
473                // Update camera matrices
474                let view_matrix = self.viewer.camera.view_matrix();
475                let proj_matrix = self.viewer.camera.projection_matrix();
476                let camera_pos = self.viewer.camera.position.coords;
477                point_renderer.update_camera(view_matrix, proj_matrix, camera_pos);
478                mesh_renderer.update_camera(view_matrix, proj_matrix, camera_pos);
479
480                // Rebuild quad vertices only when data has changed
481                if self.viewer.vertices_dirty {
482                    self.cached_vertices = match &self.viewer.current_data {
483                        ViewData::PointCloud(cloud) => {
484                            let mut vertices = Vec::with_capacity(cloud.len() * 6);
485                            for point in cloud.iter() {
486                                let size = 0.02;
487                                let pos = [point.x, point.y, point.z];
488                                let color = [1.0, 1.0, 1.0];
489                                let normal = [0.0, 0.0, 1.0];
490
491                                let v1 = PointVertex::from_point(
492                                    &Point3f::new(pos[0] - size, pos[1] - size, pos[2]),
493                                    color,
494                                    16.0,
495                                    normal,
496                                );
497                                let v2 = PointVertex::from_point(
498                                    &Point3f::new(pos[0] + size, pos[1] - size, pos[2]),
499                                    color,
500                                    16.0,
501                                    normal,
502                                );
503                                let v3 = PointVertex::from_point(
504                                    &Point3f::new(pos[0] + size, pos[1] + size, pos[2]),
505                                    color,
506                                    16.0,
507                                    normal,
508                                );
509                                let v4 = PointVertex::from_point(
510                                    &Point3f::new(pos[0] - size, pos[1] + size, pos[2]),
511                                    color,
512                                    16.0,
513                                    normal,
514                                );
515
516                                vertices.push(v1);
517                                vertices.push(v2);
518                                vertices.push(v3);
519                                vertices.push(v1);
520                                vertices.push(v3);
521                                vertices.push(v4);
522                            }
523                            vertices
524                        }
525                        ViewData::ColoredPointCloud(cloud) => {
526                            let mut vertices = Vec::with_capacity(cloud.len() * 6);
527                            for point in cloud.iter() {
528                                let size = 0.02;
529                                let pos = [point.position.x, point.position.y, point.position.z];
530                                let color = [
531                                    point.color[0] as f32 / 255.0,
532                                    point.color[1] as f32 / 255.0,
533                                    point.color[2] as f32 / 255.0,
534                                ];
535                                let normal = [0.0, 0.0, 1.0];
536
537                                let v1 = PointVertex::from_point(
538                                    &Point3f::new(pos[0] - size, pos[1] - size, pos[2]),
539                                    color,
540                                    16.0,
541                                    normal,
542                                );
543                                let v2 = PointVertex::from_point(
544                                    &Point3f::new(pos[0] + size, pos[1] - size, pos[2]),
545                                    color,
546                                    16.0,
547                                    normal,
548                                );
549                                let v3 = PointVertex::from_point(
550                                    &Point3f::new(pos[0] + size, pos[1] + size, pos[2]),
551                                    color,
552                                    16.0,
553                                    normal,
554                                );
555                                let v4 = PointVertex::from_point(
556                                    &Point3f::new(pos[0] - size, pos[1] + size, pos[2]),
557                                    color,
558                                    16.0,
559                                    normal,
560                                );
561
562                                vertices.push(v1);
563                                vertices.push(v2);
564                                vertices.push(v3);
565                                vertices.push(v1);
566                                vertices.push(v3);
567                                vertices.push(v4);
568                            }
569                            vertices
570                        }
571                        ViewData::Mesh(_) | ViewData::Empty => vec![],
572                    };
573                    self.viewer.vertices_dirty = false;
574                }
575
576                // Debug: Print vertex count periodically
577                if !self.cached_vertices.is_empty() {
578                    if self.viewer.debug_frame_count % 60 == 0 {
579                        println!("Rendering {} vertices", self.cached_vertices.len());
580                    }
581                }
582                self.viewer.debug_frame_count += 1;
583
584                match &self.viewer.current_data {
585                    ViewData::PointCloud(_) | ViewData::ColoredPointCloud(_) => {
586                        if self.screenshot_pending {
587                            self.screenshot_pending = false;
588                            println!("Screenshot not yet supported for point clouds.");
589                        }
590                        if !self.cached_vertices.is_empty() {
591                            if let Err(e) = point_renderer.render(&self.cached_vertices) {
592                                eprintln!("Render error: {}", e);
593                            }
594                        }
595                    }
596                    ViewData::Mesh(mesh) => {
597                        if !mesh.vertices.is_empty() && !mesh.faces.is_empty() {
598                            let indices: Vec<u32> = mesh
599                                .faces
600                                .iter()
601                                .flat_map(|f| [f[0] as u32, f[1] as u32, f[2] as u32])
602                                .collect();
603
604                            let normals_opt = mesh.normals.as_ref().map(|n| n.as_slice());
605
606                            let colors_f32: Option<Vec<[f32; 3]>> =
607                                mesh.colors.as_ref().map(|cols| {
608                                    cols.iter()
609                                        .map(|c| {
610                                            [
611                                                c[0] as f32 / 255.0,
612                                                c[1] as f32 / 255.0,
613                                                c[2] as f32 / 255.0,
614                                            ]
615                                        })
616                                        .collect()
617                                });
618                            let colors_opt = colors_f32.as_ref().map(|c| c.as_slice());
619
620                            // Use the viewer's material and shading mode
621                            let gpu_mesh = mesh_to_gpu_mesh(
622                                &mesh.vertices,
623                                &indices,
624                                normals_opt,
625                                colors_opt,
626                                Some(self.viewer.material),
627                            );
628
629                            if let Err(e) =
630                                mesh_renderer.render(&gpu_mesh, self.viewer.shading_mode)
631                            {
632                                eprintln!("Mesh render error: {}", e);
633                            }
634
635                            // Handle screenshot request
636                            if self.screenshot_pending {
637                                self.screenshot_pending = false;
638                                match mesh_renderer
639                                    .render_to_texture(&gpu_mesh, self.viewer.shading_mode)
640                                {
641                                    Ok((pixels, format, w, h)) => {
642                                        save_screenshot(&pixels, format, w, h);
643                                    }
644                                    Err(e) => eprintln!("Screenshot render error: {}", e),
645                                }
646                            }
647                        }
648                    }
649                    ViewData::Empty => {}
650                }
651
652                // Request next frame
653                window.request_redraw();
654            }
655            _ => {}
656        }
657    }
658
659    fn about_to_wait(&mut self, _event_loop: &ActiveEventLoop) {
660        if let Some(window) = &self.window {
661            window.request_redraw();
662        }
663    }
664}
665
666/// Encode raw GPU pixel bytes as a PNG file.
667///
668/// Handles both RGBA and BGRA surface formats transparently.
669fn save_screenshot(pixels: &[u8], format: wgpu::TextureFormat, width: u32, height: u32) {
670    use wgpu::TextureFormat;
671
672    // Convert BGRA → RGBA if the surface uses a BGRA format (common on macOS/Metal)
673    let rgba: Vec<u8> = match format {
674        TextureFormat::Bgra8Unorm | TextureFormat::Bgra8UnormSrgb => {
675            let mut buf = Vec::with_capacity(pixels.len());
676            for chunk in pixels.chunks_exact(4) {
677                buf.push(chunk[2]); // R ← B
678                buf.push(chunk[1]); // G
679                buf.push(chunk[0]); // B ← R
680                buf.push(chunk[3]); // A
681            }
682            buf
683        }
684        _ => pixels.to_vec(),
685    };
686
687    let timestamp = std::time::SystemTime::now()
688        .duration_since(std::time::UNIX_EPOCH)
689        .map(|d| d.as_secs())
690        .unwrap_or(0);
691    let path = format!("screenshot_{}.png", timestamp);
692
693    match image::RgbaImage::from_raw(width, height, rgba) {
694        Some(img) => match img.save(&path) {
695            Ok(()) => println!("Screenshot saved: {}", path),
696            Err(e) => eprintln!("Failed to save screenshot: {}", e),
697        },
698        None => eprintln!("Failed to build image buffer for screenshot"),
699    }
700}