1use 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#[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 #[arg(short, long, num_args = 1..)]
41 models: Vec<String>,
42
43 #[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#[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 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 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}