Skip to main content

simple/
simple.rs

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