1use 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#[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 #[arg(short, long)]
42 model: String,
43
44 #[arg(long, default_value_t = 1.0)]
46 size: f32,
47
48 #[arg(long, value_enum, default_value_t = DisplayMode::Splat, ignore_case = true)]
50 mode: DisplayMode,
51
52 #[arg(long, default_value_t = 3, value_parser = clap::value_parser!(u8).range(0..=3))]
54 sh_degree: u8,
55
56 #[arg(long, default_value_t)]
58 no_sh0: bool,
59
60 #[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#[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 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 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 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}