wgpu_3dgs_viewer/
lib.rs

1#![doc = include_str!("../README.md")]
2
3mod buffer;
4mod camera;
5mod error;
6mod preprocessor;
7mod radix_sorter;
8mod renderer;
9pub mod shader;
10mod wesl_utils;
11
12#[cfg(feature = "multi-model")]
13mod multi_model;
14
15#[cfg(feature = "selection")]
16pub mod selection;
17
18use glam::*;
19use wgpu_3dgs_core::{
20    BufferWrapper, GaussianDisplayMode, GaussianMaxStdDev, GaussianPod, GaussianShDegree,
21    GaussianTransformBuffer, GaussianTransformPod, GaussiansBuffer, IterGaussian,
22    ModelTransformBuffer, ModelTransformPod,
23};
24
25#[cfg(feature = "viewer-selection")]
26use wgpu_3dgs_editor::SelectionBuffer;
27
28pub use buffer::*;
29pub use camera::*;
30pub use error::*;
31pub use preprocessor::*;
32pub use radix_sorter::*;
33pub use renderer::*;
34
35#[cfg(feature = "multi-model")]
36pub use multi_model::*;
37
38pub use wgpu_3dgs_core as core;
39
40#[cfg(feature = "editor")]
41pub use wgpu_3dgs_editor as editor;
42
43/// The default viewer [`GaussianPod`] type.
44pub type DefaultGaussianPod = core::GaussianPodWithShSingleCov3dSingleConfigs;
45
46/// The 3D Gaussian splatting viewer.
47///
48/// This provides all the necessary buffers and operations to render 3D Gaussians:
49/// - Buffers
50///     - [`CameraBuffer`]
51///     - [`ModelTransformBuffer`]
52///     - [`GaussianTransformBuffer`]
53///     - [`GaussiansBuffer`]
54///     - [`IndirectArgsBuffer`]
55///     - [`RadixSortIndirectArgsBuffer`]
56///     - [`IndirectIndicesBuffer`]
57///     - [`GaussiansDepthBuffer`]
58/// - Operations
59///     - [`Preprocessor`]
60///     - [`RadixSorter`]
61///     - [`Renderer`]
62///
63/// If you wish to manage these buffers yourself, you do not need to use this struct.
64#[derive(Debug)]
65pub struct Viewer<G: GaussianPod = DefaultGaussianPod> {
66    pub camera_buffer: CameraBuffer,
67    pub model_transform_buffer: ModelTransformBuffer,
68    pub gaussian_transform_buffer: GaussianTransformBuffer,
69    pub gaussians_buffer: GaussiansBuffer<G>,
70    pub indirect_args_buffer: IndirectArgsBuffer,
71    pub radix_sort_indirect_args_buffer: RadixSortIndirectArgsBuffer,
72    pub indirect_indices_buffer: IndirectIndicesBuffer,
73    pub gaussians_depth_buffer: GaussiansDepthBuffer,
74    #[cfg(feature = "viewer-selection")]
75    pub selection_buffer: SelectionBuffer,
76    #[cfg(feature = "viewer-selection")]
77    pub invert_selection_buffer: selection::PreprocessorInvertSelectionBuffer,
78
79    pub preprocessor: Preprocessor<G>,
80    pub radix_sorter: RadixSorter,
81    pub renderer: Renderer<G>,
82}
83
84impl<G: GaussianPod> Viewer<G> {
85    /// Create a new viewer.
86    pub fn new(
87        device: &wgpu::Device,
88        texture_format: wgpu::TextureFormat,
89        gaussians: &impl IterGaussian,
90    ) -> Result<Self, ViewerCreateError> {
91        Self::new_with_options(
92            device,
93            texture_format,
94            gaussians,
95            ViewerCreateOptions::default(),
96        )
97    }
98
99    /// Create a new viewer with extra [`ViewerCreateOptions`].
100    pub fn new_with_options(
101        device: &wgpu::Device,
102        texture_format: wgpu::TextureFormat,
103        gaussians: &impl IterGaussian,
104        options: ViewerCreateOptions,
105    ) -> Result<Self, ViewerCreateError> {
106        log::debug!("Creating camera buffer");
107        let camera_buffer = CameraBuffer::new(device);
108
109        log::debug!("Creating model transform buffer");
110        let model_transform_buffer = ModelTransformBuffer::new(device);
111
112        log::debug!("Creating gaussian transform buffer");
113        let gaussian_transform_buffer = GaussianTransformBuffer::new(device);
114
115        log::debug!("Creating gaussians buffer");
116        let gaussians_buffer =
117            GaussiansBuffer::new_with_usage(device, gaussians, options.gaussians_buffer_usage);
118
119        log::debug!("Creating indirect args buffer");
120        let indirect_args_buffer = IndirectArgsBuffer::new(device);
121
122        log::debug!("Creating radix sort indirect args buffer");
123        let radix_sort_indirect_args_buffer = RadixSortIndirectArgsBuffer::new(device);
124
125        // Assuming it is cheap to call `iter_gaussian`.
126        let len = gaussians.iter_gaussian().len() as u32;
127
128        log::debug!("Creating indirect indices buffer");
129        let indirect_indices_buffer = IndirectIndicesBuffer::new(device, len);
130
131        log::debug!("Creating gaussians depth buffer");
132        let gaussians_depth_buffer = GaussiansDepthBuffer::new(device, len);
133
134        #[cfg(feature = "viewer-selection")]
135        let selection_buffer = {
136            log::debug!("Creating selection buffer");
137            SelectionBuffer::new(device, len)
138        };
139
140        #[cfg(feature = "viewer-selection")]
141        let invert_selection_buffer = {
142            log::debug!("Creating invert selection buffer");
143            selection::PreprocessorInvertSelectionBuffer::new(device)
144        };
145
146        log::debug!("Creating preprocessor");
147        let preprocessor = Preprocessor::new(
148            device,
149            &camera_buffer,
150            &model_transform_buffer,
151            &gaussian_transform_buffer,
152            &gaussians_buffer,
153            &indirect_args_buffer,
154            &radix_sort_indirect_args_buffer,
155            &indirect_indices_buffer,
156            &gaussians_depth_buffer,
157            #[cfg(feature = "viewer-selection")]
158            &selection_buffer,
159            #[cfg(feature = "viewer-selection")]
160            &invert_selection_buffer,
161        )?;
162
163        log::debug!("Creating radix sorter");
164        let radix_sorter =
165            RadixSorter::new(device, &gaussians_depth_buffer, &indirect_indices_buffer);
166
167        log::debug!("Creating renderer");
168        let renderer = Renderer::new(
169            device,
170            texture_format,
171            options.depth_stencil,
172            &camera_buffer,
173            &model_transform_buffer,
174            &gaussian_transform_buffer,
175            &gaussians_buffer,
176            &indirect_indices_buffer,
177        )?;
178
179        log::info!("Viewer created");
180
181        Ok(Self {
182            camera_buffer,
183            model_transform_buffer,
184            gaussian_transform_buffer,
185            gaussians_buffer,
186            indirect_args_buffer,
187            radix_sort_indirect_args_buffer,
188            indirect_indices_buffer,
189            gaussians_depth_buffer,
190            #[cfg(feature = "viewer-selection")]
191            selection_buffer,
192            #[cfg(feature = "viewer-selection")]
193            invert_selection_buffer,
194
195            preprocessor,
196            radix_sorter,
197            renderer,
198        })
199    }
200
201    /// Update the camera.
202    pub fn update_camera(
203        &mut self,
204        queue: &wgpu::Queue,
205        camera: &impl CameraTrait,
206        texture_size: UVec2,
207    ) {
208        self.camera_buffer.update(queue, camera, texture_size);
209    }
210
211    /// Update the camera with [`CameraPod`].
212    pub fn update_camera_with_pod(&mut self, queue: &wgpu::Queue, pod: &CameraPod) {
213        self.camera_buffer.update_with_pod(queue, pod);
214    }
215
216    /// Update the model transform.
217    pub fn update_model_transform(
218        &mut self,
219        queue: &wgpu::Queue,
220        pos: Vec3,
221        rot: Quat,
222        scale: Vec3,
223    ) {
224        self.model_transform_buffer.update(queue, pos, rot, scale);
225    }
226
227    /// Update the model transform with [`ModelTransformPod`].
228    pub fn update_model_transform_with_pod(
229        &mut self,
230        queue: &wgpu::Queue,
231        pod: &ModelTransformPod,
232    ) {
233        self.model_transform_buffer.update_with_pod(queue, pod);
234    }
235
236    /// Update the Gaussian transform.
237    pub fn update_gaussian_transform(
238        &mut self,
239        queue: &wgpu::Queue,
240        size: f32,
241        display_mode: GaussianDisplayMode,
242        sh_deg: GaussianShDegree,
243        no_sh0: bool,
244        max_std_dev: GaussianMaxStdDev,
245    ) {
246        self.gaussian_transform_buffer.update(
247            queue,
248            size,
249            display_mode,
250            sh_deg,
251            no_sh0,
252            max_std_dev,
253        );
254    }
255
256    /// Update the Gaussian transform with [`GaussianTransformPod`].
257    pub fn update_gaussian_transform_with_pod(
258        &mut self,
259        queue: &wgpu::Queue,
260        pod: &GaussianTransformPod,
261    ) {
262        self.gaussian_transform_buffer.update_with_pod(queue, pod);
263    }
264
265    /// Render the viewer.
266    pub fn render(&self, encoder: &mut wgpu::CommandEncoder, texture_view: &wgpu::TextureView) {
267        self.preprocessor
268            .preprocess(encoder, self.gaussians_buffer.len() as u32);
269
270        self.radix_sorter
271            .sort(encoder, &self.radix_sort_indirect_args_buffer);
272
273        self.renderer
274            .render(encoder, texture_view, &self.indirect_args_buffer);
275    }
276}
277
278/// The options for creating a [`Viewer`] using [`Viewer::new_with_options`].
279pub struct ViewerCreateOptions {
280    /// The optional depth stencil state for the renderer.
281    pub depth_stencil: Option<wgpu::DepthStencilState>,
282    /// The usage for the gaussians buffer.
283    pub gaussians_buffer_usage: wgpu::BufferUsages,
284}
285
286impl Default for ViewerCreateOptions {
287    fn default() -> Self {
288        Self {
289            depth_stencil: None,
290            gaussians_buffer_usage: GaussiansBuffer::<DefaultGaussianPod>::DEFAULT_USAGES,
291        }
292    }
293}
wgpu_3dgs_viewer/lib.rs

wgpu_3dgs_viewer/
lib.rs
