Skip to main content

multi_model/
multi_model.rs

1//! This example renders multiple models in the order of their centroids' distance to the camera using the `multi-model` feature.
2//!
3//! For example, to use an offset of (10, 0, 0) between each model, run:
4//!
5//! ```sh
6//! cargo run --example multi-model --features="multi-model" -- -m "path/to/model1.ply" -m "path/to/model2.ply" --offset 10.0,0.0,0.0
7//! ```
8//!
9//! To view more options and the controls, run with `--help`:
10//!
11//! ```sh
12//! cargo run --example multi-model --features="multi-model" -- --help
13//! ```
14
15use std::sync::Arc;
16
17use clap::Parser;
18use glam::*;
19use winit::{error::EventLoopError, event_loop::EventLoop, keyboard::KeyCode, window::Window};
20
21use wgpu_3dgs_viewer as gs;
22use wgpu_3dgs_viewer::core::{GaussiansSource, IterGaussian};
23
24mod utils;
25use utils::core;
26
27/// The command line arguments.
28#[derive(Parser, Debug)]
29#[command(
30    version,
31    about,
32    long_about = "\
33    A 3D Gaussian splatting viewer written in Rust using wgpu.\n\
34    \n\
35    Use W, A, S, D, Space, Shift to move, use mouse to rotate.\n\
36    "
37)]
38struct Args {
39    /// Path to the .ply file.
40    #[arg(short, long, num_args = 1..)]
41    models: Vec<String>,
42
43    /// The offset of each model.
44    #[arg(
45        short,
46        long,
47        num_args = 3,
48        value_delimiter = ',',
49        default_value = "10.0,0.0,0.0"
50    )]
51    offset: Vec<f32>,
52}
53
54fn main() -> Result<(), EventLoopError> {
55    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
56
57    let event_loop = EventLoop::new()?;
58    event_loop.run_app(&mut core::App::<System>::new(Args::parse()))?;
59    Ok(())
60}
61
62/// The application system.
63#[allow(dead_code)]
64struct System {
65    surface: wgpu::Surface<'static>,
66    queue: wgpu::Queue,
67    device: wgpu::Device,
68    config: wgpu::SurfaceConfiguration,
69
70    camera: gs::Camera,
71    gaussians: Vec<gs::core::Gaussians>,
72    gaussian_centroids: Vec<Vec3>,
73    viewer: gs::MultiModelViewer<gs::DefaultGaussianPod, usize>,
74}
75
76impl core::System for System {
77    type Args = Args;
78
79    async fn init(window: Arc<Window>, args: &Args) -> Self {
80        let model_paths = &args.models;
81        let model_offset = Vec3::from_slice(&args.offset);
82        let size = window.inner_size();
83
84        log::debug!("Creating wgpu instance");
85        let instance =
86            wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle_from_env());
87
88        log::debug!("Creating window surface");
89        let surface = instance.create_surface(window.clone()).expect("surface");
90
91        log::debug!("Requesting adapter");
92        let adapter = instance
93            .request_adapter(&wgpu::RequestAdapterOptions {
94                power_preference: wgpu::PowerPreference::HighPerformance,
95                compatible_surface: Some(&surface),
96                force_fallback_adapter: false,
97            })
98            .await
99            .expect("adapter");
100
101        log::debug!("Requesting device");
102        let (device, queue) = adapter
103            .request_device(&wgpu::DeviceDescriptor {
104                label: Some("Device"),
105                required_limits: adapter.limits(),
106                ..Default::default()
107            })
108            .await
109            .expect("device");
110
111        let surface_caps = surface.get_capabilities(&adapter);
112        let surface_format = surface_caps.formats[0];
113        let config = wgpu::SurfaceConfiguration {
114            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
115            format: surface_format,
116            width: size.width.max(1),
117            height: size.height.max(1),
118            present_mode: surface_caps.present_modes[0],
119            alpha_mode: surface_caps.alpha_modes[0],
120            view_formats: vec![surface_format.remove_srgb_suffix()],
121            desired_maximum_frame_latency: 2,
122        };
123
124        log::debug!("Configuring surface");
125        surface.configure(&device, &config);
126
127        log::debug!("Creating gaussians");
128        let gaussians = model_paths
129            .iter()
130            .map(|model_path| {
131                log::debug!("Reading model from {model_path}");
132                [GaussiansSource::Ply, GaussiansSource::Spz]
133                    .into_iter()
134                    .find_map(|source| gs::core::Gaussians::read_from_file(model_path, source).ok())
135                    .expect("gaussians")
136            })
137            .collect::<Vec<_>>();
138
139        log::debug!("Computing gaussian centroids");
140        let mut gaussian_centroids = gaussians
141            .iter()
142            .map(|g| {
143                let mut centroid = Vec3::ZERO;
144                for gaussian in g.iter_gaussian() {
145                    centroid += gaussian.pos;
146                }
147                centroid / g.len() as f32
148            })
149            .collect::<Vec<_>>();
150
151        log::debug!("Creating camera");
152        let camera = gs::Camera::new(0.1..1e4, 60f32.to_radians());
153
154        log::debug!("Creating viewer");
155        let mut viewer =
156            gs::MultiModelViewer::new(&device, config.view_formats[0]).expect("viewer");
157
158        let quat = Quat::from_axis_angle(Vec3::Z, 180f32.to_radians());
159        for (i, gaussians) in gaussians.iter().enumerate() {
160            let offset = model_offset * i as f32;
161
162            log::debug!("Pushing model {i}");
163
164            viewer.insert_model(&device, i, gaussians);
165            viewer
166                .update_model_transform(&queue, &i, offset, quat, Vec3::ONE)
167                .expect("update model");
168
169            gaussian_centroids[i] = quat.mul_vec3(gaussian_centroids[i]) + offset;
170        }
171
172        log::info!("System initialized");
173
174        Self {
175            surface,
176            device,
177            queue,
178            config,
179
180            camera,
181            gaussians,
182            gaussian_centroids,
183            viewer,
184        }
185    }
186
187    fn update(&mut self, input: &core::Input, delta_time: f32) {
188        const SPEED: f32 = 1.0;
189
190        let mut forward = 0.0;
191        if input.held_keys.contains(&KeyCode::KeyW) {
192            forward += SPEED * delta_time;
193        }
194        if input.held_keys.contains(&KeyCode::KeyS) {
195            forward -= SPEED * delta_time;
196        }
197
198        let mut right = 0.0;
199        if input.held_keys.contains(&KeyCode::KeyD) {
200            right += SPEED * delta_time;
201        }
202        if input.held_keys.contains(&KeyCode::KeyA) {
203            right -= SPEED * delta_time;
204        }
205
206        self.camera.move_by(forward, right);
207
208        let mut up = 0.0;
209        if input.held_keys.contains(&KeyCode::Space) {
210            up += SPEED * delta_time;
211        }
212        if input.held_keys.contains(&KeyCode::ShiftLeft) {
213            up -= SPEED * delta_time;
214        }
215
216        self.camera.move_up(up);
217
218        // Camera rotation
219        const SENSITIVITY: f32 = 0.15;
220
221        let yaw = input.mouse_diff.x * SENSITIVITY * delta_time;
222        let pitch = input.mouse_diff.y * SENSITIVITY * delta_time;
223
224        self.camera.pitch_by(-pitch);
225        self.camera.yaw_by(-yaw);
226
227        // Update the viewer
228        self.viewer.update_camera(
229            &self.queue,
230            &self.camera,
231            uvec2(self.config.width, self.config.height),
232        );
233    }
234
235    fn render(&mut self) {
236        let texture = match self.surface.get_current_texture() {
237            wgpu::CurrentSurfaceTexture::Success(texture)
238            | wgpu::CurrentSurfaceTexture::Suboptimal(texture) => texture,
239            e => {
240                log::error!("Failed to get current texture: {e:?}");
241                return;
242            }
243        };
244        let texture_view = texture.texture.create_view(&wgpu::TextureViewDescriptor {
245            label: Some("Texture View"),
246            format: Some(self.config.view_formats[0]),
247            ..Default::default()
248        });
249
250        let mut encoder = self
251            .device
252            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
253                label: Some("Command Encoder"),
254            });
255
256        let mut render_keys = self
257            .gaussian_centroids
258            .iter()
259            .enumerate()
260            .map(|(i, centroid)| (i, centroid - self.camera.pos))
261            .collect::<Vec<_>>();
262
263        render_keys.sort_by(|(_, a), (_, b)| {
264            a.length()
265                .partial_cmp(&b.length())
266                .unwrap_or(std::cmp::Ordering::Equal)
267        });
268
269        let render_keys = render_keys
270            .into_iter()
271            .rev()
272            .map(|(i, _)| i)
273            .collect::<Vec<_>>();
274
275        self.viewer
276            .render(
277                &mut encoder,
278                &texture_view,
279                render_keys.iter().collect::<Vec<_>>().as_slice(),
280            )
281            .expect("render");
282
283        self.queue.submit(std::iter::once(encoder.finish()));
284        if let Err(e) = self.device.poll(wgpu::PollType::wait_indefinitely()) {
285            log::error!("Failed to poll device: {e:?}");
286        }
287        texture.present();
288    }
289
290    fn resize(&mut self, size: winit::dpi::PhysicalSize<u32>) {
291        if size.width > 0 && size.height > 0 {
292            self.config.width = size.width;
293            self.config.height = size.height;
294            self.surface.configure(&self.device, &self.config);
295        }
296    }
297}