pub struct SelectionBundle<G: GaussianPod> {
pub bundles: Vec<ComputeBundle<()>>,
/* private fields */
}Expand description
A collection of specialized ComputeBundle for selection operations.
§Custom Operations
All ComputeBundles supplied to this bundle as a SelectionExpr::Unary,
SelectionExpr::Binary, or SelectionExpr::Selection custom operation must have the same
bind group 0 as the SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR. They must also
not have the bind group itself, as it will be supplied automatically during evaluation.
It is recommended to use ComputeBundleBuilder to create the custom operation bundles,
and build them using ComputeBundleBuilder::build_without_bind_groups.
// Create the selection custom operation compute bundle
let my_selection_custom_op_bundle = ComputeBundleBuilder::new()
.label("My Selection")
.bind_group_layouts([
&SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR,
&MY_CUSTOM_BIND_GROUP_LAYOUT_DESCRIPTOR, // Put your custom bind group layout here.
])
.resolver({
let mut resolver = wesl::PkgResolver::new();
resolver.add_package(&core::shader::PACKAGE); // Required for using core buffer structs.
resolver.add_package(&shader::PACKAGE); // Optionally add this for some utility functions.
resolver
})
.main_shader("path::to::my::wesl::module".parse().unwrap())
.entry_point("main")
.wesl_compile_options(wesl::CompileOptions {
features: G::wesl_features(), // Required for enabling the correct features for core structs.
..Default::default()
})
.build_without_bind_groups(&device)
.map_err(|e| log::error!("{e}"))
.expect("my selection custom op bundle");
// Create the selection bundle
let selection_bundle = SelectionBundle::<GaussianPod>::new(
&device,
vec![my_selection_custom_op_bundle],
);
// Create the bind group for your custom operation
let my_selection_custom_op_bind_group = selection_bundle.bundles[0]
.create_bind_group(
&device,
1, // Index 0 is the Gaussians buffer, so remember to start from 1 for your bind groups
[my_selection_custom_op_buffer.buffer().as_entire_binding()], // Your custom bind group resources
)
.unwrap()
// Create the selection expression
let selection_expr = gs::SelectionExpr::selection(
0, // The bundle index for your custom operation in the selection bundle
vec![my_selection_custom_op_bind_group],
)
.union( // Combine with other selection expressions using different functions
gs::SelectionExpr::Buffer(my_existing_selection_buffer) // An existing selection buffer for example
);
// Create a selection buffer for the result
let dest_selection_buffer = SelectionBuffer::new(&device, gaussians_buffer.len() as u32);
// Evaluate the selection expression
selection_bundle.evaluate(
&device,
&mut encoder,
&selection_expr,
&dest_selection_buffer,
&model_transform,
&gaussian_transform,
&gaussians_buffer,
);§Shader Format
You may copy and paste the following shader bindings for
SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR into your custom selection operation
shader to ensure that the bindings are correct, then add your own bindings after that.
import wgpu_3dgs_core::{
gaussian::Gaussian,
gaussian_transform::GaussianTransform,
model_transform::{model_to_world, ModelTransform},
};
@group(0) @binding(0)
var<uniform> op: u32;
@group(0) @binding(1)
var<storage, read> source: array<u32>;
@group(0) @binding(2)
var<storage, read_write> dest: array<atomic<u32>>;
@group(0) @binding(3)
var<uniform> model_transform: ModelTransform;
@group(0) @binding(4)
var<uniform> gaussian_transform: GaussianTransform;
@group(0) @binding(5)
var<storage, read> gaussians: array<Gaussian>;
// Your custom bindings here...
override workgroup_size: u32;
@compute @workgroup_size(workgroup_size)
fn main(@builtin(global_invocation_id) id: vec3<u32>) {
let index = id.x;
if index >= arrayLength(&gaussians) {
return;
}
let gaussian = gaussians[index];
let world_pos = model_to_world(model_transform, gaussian.position);
// Your custom selection operation code here...
let word_index = index / 32u;
let bit_index = index % 32u;
let bit_mask = 1u << bit_index;
if /* Condition for selecting the Gaussian */ {
atomicOr(&dest[word_index], bit_mask);
} else {
atomicAnd(&dest[word_index], ~bit_mask);
}
}Fields§
§bundles: Vec<ComputeBundle<()>>The compute bundles for selection custom operations.
Implementations§
Source§impl<G: GaussianPod> SelectionBundle<G>
impl<G: GaussianPod> SelectionBundle<G>
Sourcepub const GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR: BindGroupLayoutDescriptor<'static>
pub const GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR: BindGroupLayoutDescriptor<'static>
The Gaussians bind group layout descriptors.
This bind group layout takes the following buffers:
Sourcepub const SPHERE_BIND_GROUP_LAYOUT_DESCRIPTOR: BindGroupLayoutDescriptor<'static>
pub const SPHERE_BIND_GROUP_LAYOUT_DESCRIPTOR: BindGroupLayoutDescriptor<'static>
The sphere selection bind group layout descriptor.
This bind group layout takes the following buffers:
Sourcepub const BOX_BIND_GROUP_LAYOUT_DESCRIPTOR: BindGroupLayoutDescriptor<'static>
pub const BOX_BIND_GROUP_LAYOUT_DESCRIPTOR: BindGroupLayoutDescriptor<'static>
The box selection bind group layout descriptor.
This bind group layout takes the following buffers:
Sourcepub fn new<'a>(device: &Device, bundles: Vec<ComputeBundle<()>>) -> Self
pub fn new<'a>(device: &Device, bundles: Vec<ComputeBundle<()>>) -> Self
Create a new selection bundle.
bundles are used for SelectionExpr::Unary, SelectionExpr::Binary, or
SelectionExpr::Selection as custom operations, they must have the same bind group 0 as
the SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR, see documentation of
SelectionBundle for more details.
Examples found in repository?
88async fn main() {
89 env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
90
91 let args = Args::parse();
92 let model_path = &args.model;
93 let pos = Vec3::from_slice(&args.pos);
94 let rot = Quat::from_slice(&args.rot);
95 let scale = Vec3::from_slice(&args.scale);
96 let shape = match args.shape {
97 Shape::Sphere => gs::SelectionBundle::<GaussianPod>::create_sphere_bundle,
98 Shape::Box => gs::SelectionBundle::<GaussianPod>::create_box_bundle,
99 };
100 let repeat = args.repeat;
101 let offset = Vec3::from_slice(&args.offset);
102
103 log::debug!("Creating wgpu instance");
104 let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
105
106 log::debug!("Requesting adapter");
107 let adapter = instance
108 .request_adapter(&wgpu::RequestAdapterOptions::default())
109 .await
110 .expect("adapter");
111
112 log::debug!("Requesting device");
113 let (device, queue) = adapter
114 .request_device(&wgpu::DeviceDescriptor {
115 label: Some("Device"),
116 required_features: wgpu::Features::empty(),
117 required_limits: adapter.limits(),
118 memory_hints: wgpu::MemoryHints::default(),
119 trace: wgpu::Trace::Off,
120 })
121 .await
122 .expect("device");
123
124 log::debug!("Creating gaussians");
125 let f = std::fs::File::open(model_path).expect("ply file");
126 let mut reader = std::io::BufReader::new(f);
127 let gaussians = gs::core::Gaussians::read_ply(&mut reader).expect("gaussians");
128
129 log::debug!("Creating gaussians buffer");
130 let gaussians_buffer =
131 gs::core::GaussiansBuffer::<GaussianPod>::new(&device, &gaussians.gaussians);
132
133 log::debug!("Creating model transform buffer");
134 let model_transform = gs::core::ModelTransformBuffer::new(&device);
135
136 log::debug!("Creating Gaussian transform buffer");
137 let gaussian_transform = gs::core::GaussianTransformBuffer::new(&device);
138
139 log::debug!("Creating shape selection compute bundle");
140 let shape_selection = shape(&device);
141
142 log::debug!("Creating selection bundle");
143 let selection_bundle = gs::SelectionBundle::<GaussianPod>::new(&device, vec![shape_selection]);
144
145 log::debug!("Creating shape selection buffers");
146 let shape_selection_buffers = (0..repeat)
147 .map(|i| {
148 let offset_pos = pos + offset * i as f32;
149 let buffer = gs::InvTransformBuffer::new(&device);
150 buffer.update_with_scale_rot_pos(&queue, scale, rot, offset_pos);
151 buffer
152 })
153 .collect::<Vec<_>>();
154
155 log::debug!("Creating shape selection bind groups");
156 let shape_selection_bind_groups = shape_selection_buffers
157 .iter()
158 .map(|buffer| {
159 selection_bundle.bundles[0]
160 .create_bind_group(
161 &device,
162 // index 0 is the Gaussians buffer, so we use 1,
163 // see docs of create_sphere_bundle or create_box_bundle
164 1,
165 [buffer.buffer().as_entire_binding()],
166 )
167 .expect("bind group")
168 })
169 .collect::<Vec<_>>();
170
171 log::debug!("Creating selection expression");
172 let selection_expr = shape_selection_bind_groups.into_iter().fold(
173 gs::SelectionExpr::Identity,
174 |acc, bind_group| {
175 acc.union(gs::SelectionExpr::selection(
176 0, // the 0 here is the bundle index in the selection bundle
177 vec![bind_group],
178 ))
179 },
180 );
181
182 log::debug!("Creating destination buffer");
183 let dest = gs::SelectionBuffer::new(&device, gaussians_buffer.len() as u32);
184
185 log::info!("Starting selection process");
186 let time = std::time::Instant::now();
187
188 log::debug!("Selecting Gaussians");
189 let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
190 label: Some("Selection Encoder"),
191 });
192
193 selection_bundle.evaluate(
194 &device,
195 &mut encoder,
196 &selection_expr,
197 &dest,
198 &model_transform,
199 &gaussian_transform,
200 &gaussians_buffer,
201 );
202
203 queue.submit(Some(encoder.finish()));
204
205 #[allow(unused_must_use)]
206 device.poll(wgpu::PollType::Wait);
207
208 log::info!("Editing process completed in {:?}", time.elapsed());
209
210 log::debug!("Filtering Gaussians");
211 let selected_gaussians = gs::core::Gaussians {
212 gaussians: dest
213 .download::<u32>(&device, &queue)
214 .await
215 .expect("selected download")
216 .iter()
217 .flat_map(|group| {
218 std::iter::repeat_n(group, 32)
219 .enumerate()
220 .map(|(i, g)| g & (1 << i) != 0)
221 })
222 .zip(gaussians.gaussians.iter())
223 .filter(|(selected, _)| *selected)
224 .map(|(_, g)| g.clone())
225 .collect::<Vec<_>>(),
226 };
227
228 log::debug!("Writing modified Gaussians to output file");
229 let output_file = std::fs::File::create(&args.output).expect("output file");
230 let mut writer = std::io::BufWriter::new(output_file);
231 selected_gaussians
232 .write_ply(&mut writer)
233 .expect("write modified Gaussians to output file");
234
235 log::info!("Filtered Gaussians written to {}", args.output);
236}Sourcepub fn gaussians_bind_group_layout(&self) -> &BindGroupLayout
pub fn gaussians_bind_group_layout(&self) -> &BindGroupLayout
Get the Gaussians bind group layout.
Sourcepub fn evaluate(
&self,
device: &Device,
encoder: &mut CommandEncoder,
expr: &SelectionExpr,
dest: &SelectionBuffer,
model_transform: &ModelTransformBuffer,
gaussian_transform: &GaussianTransformBuffer,
gaussians: &GaussiansBuffer<G>,
)
pub fn evaluate( &self, device: &Device, encoder: &mut CommandEncoder, expr: &SelectionExpr, dest: &SelectionBuffer, model_transform: &ModelTransformBuffer, gaussian_transform: &GaussianTransformBuffer, gaussians: &GaussiansBuffer<G>, )
Evaluate and apply the selection expression.
Examples found in repository?
88async fn main() {
89 env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
90
91 let args = Args::parse();
92 let model_path = &args.model;
93 let pos = Vec3::from_slice(&args.pos);
94 let rot = Quat::from_slice(&args.rot);
95 let scale = Vec3::from_slice(&args.scale);
96 let shape = match args.shape {
97 Shape::Sphere => gs::SelectionBundle::<GaussianPod>::create_sphere_bundle,
98 Shape::Box => gs::SelectionBundle::<GaussianPod>::create_box_bundle,
99 };
100 let repeat = args.repeat;
101 let offset = Vec3::from_slice(&args.offset);
102
103 log::debug!("Creating wgpu instance");
104 let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
105
106 log::debug!("Requesting adapter");
107 let adapter = instance
108 .request_adapter(&wgpu::RequestAdapterOptions::default())
109 .await
110 .expect("adapter");
111
112 log::debug!("Requesting device");
113 let (device, queue) = adapter
114 .request_device(&wgpu::DeviceDescriptor {
115 label: Some("Device"),
116 required_features: wgpu::Features::empty(),
117 required_limits: adapter.limits(),
118 memory_hints: wgpu::MemoryHints::default(),
119 trace: wgpu::Trace::Off,
120 })
121 .await
122 .expect("device");
123
124 log::debug!("Creating gaussians");
125 let f = std::fs::File::open(model_path).expect("ply file");
126 let mut reader = std::io::BufReader::new(f);
127 let gaussians = gs::core::Gaussians::read_ply(&mut reader).expect("gaussians");
128
129 log::debug!("Creating gaussians buffer");
130 let gaussians_buffer =
131 gs::core::GaussiansBuffer::<GaussianPod>::new(&device, &gaussians.gaussians);
132
133 log::debug!("Creating model transform buffer");
134 let model_transform = gs::core::ModelTransformBuffer::new(&device);
135
136 log::debug!("Creating Gaussian transform buffer");
137 let gaussian_transform = gs::core::GaussianTransformBuffer::new(&device);
138
139 log::debug!("Creating shape selection compute bundle");
140 let shape_selection = shape(&device);
141
142 log::debug!("Creating selection bundle");
143 let selection_bundle = gs::SelectionBundle::<GaussianPod>::new(&device, vec![shape_selection]);
144
145 log::debug!("Creating shape selection buffers");
146 let shape_selection_buffers = (0..repeat)
147 .map(|i| {
148 let offset_pos = pos + offset * i as f32;
149 let buffer = gs::InvTransformBuffer::new(&device);
150 buffer.update_with_scale_rot_pos(&queue, scale, rot, offset_pos);
151 buffer
152 })
153 .collect::<Vec<_>>();
154
155 log::debug!("Creating shape selection bind groups");
156 let shape_selection_bind_groups = shape_selection_buffers
157 .iter()
158 .map(|buffer| {
159 selection_bundle.bundles[0]
160 .create_bind_group(
161 &device,
162 // index 0 is the Gaussians buffer, so we use 1,
163 // see docs of create_sphere_bundle or create_box_bundle
164 1,
165 [buffer.buffer().as_entire_binding()],
166 )
167 .expect("bind group")
168 })
169 .collect::<Vec<_>>();
170
171 log::debug!("Creating selection expression");
172 let selection_expr = shape_selection_bind_groups.into_iter().fold(
173 gs::SelectionExpr::Identity,
174 |acc, bind_group| {
175 acc.union(gs::SelectionExpr::selection(
176 0, // the 0 here is the bundle index in the selection bundle
177 vec![bind_group],
178 ))
179 },
180 );
181
182 log::debug!("Creating destination buffer");
183 let dest = gs::SelectionBuffer::new(&device, gaussians_buffer.len() as u32);
184
185 log::info!("Starting selection process");
186 let time = std::time::Instant::now();
187
188 log::debug!("Selecting Gaussians");
189 let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
190 label: Some("Selection Encoder"),
191 });
192
193 selection_bundle.evaluate(
194 &device,
195 &mut encoder,
196 &selection_expr,
197 &dest,
198 &model_transform,
199 &gaussian_transform,
200 &gaussians_buffer,
201 );
202
203 queue.submit(Some(encoder.finish()));
204
205 #[allow(unused_must_use)]
206 device.poll(wgpu::PollType::Wait);
207
208 log::info!("Editing process completed in {:?}", time.elapsed());
209
210 log::debug!("Filtering Gaussians");
211 let selected_gaussians = gs::core::Gaussians {
212 gaussians: dest
213 .download::<u32>(&device, &queue)
214 .await
215 .expect("selected download")
216 .iter()
217 .flat_map(|group| {
218 std::iter::repeat_n(group, 32)
219 .enumerate()
220 .map(|(i, g)| g & (1 << i) != 0)
221 })
222 .zip(gaussians.gaussians.iter())
223 .filter(|(selected, _)| *selected)
224 .map(|(_, g)| g.clone())
225 .collect::<Vec<_>>(),
226 };
227
228 log::debug!("Writing modified Gaussians to output file");
229 let output_file = std::fs::File::create(&args.output).expect("output file");
230 let mut writer = std::io::BufWriter::new(output_file);
231 selected_gaussians
232 .write_ply(&mut writer)
233 .expect("write modified Gaussians to output file");
234
235 log::info!("Filtered Gaussians written to {}", args.output);
236}Sourcepub fn create_primitive_bundle(device: &Device) -> ComputeBundle<()>
pub fn create_primitive_bundle(device: &Device) -> ComputeBundle<()>
Create the selection primitive operation ComputeBundle.
- Bind group 0 is
SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR.
You usually do not need to use this method, it is used internally for creating the primitive operation bundle for evaluation.
Sourcepub fn create_sphere_bundle(device: &Device) -> ComputeBundle<()>
pub fn create_sphere_bundle(device: &Device) -> ComputeBundle<()>
Create a sphere selection custom operation.
- Bind group 0 is
SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR. - Bind group 1 is
SelectionBundle::SPHERE_BIND_GROUP_LAYOUT_DESCRIPTOR.
Sourcepub fn create_box_bundle(device: &Device) -> ComputeBundle<()>
pub fn create_box_bundle(device: &Device) -> ComputeBundle<()>
Create a box selection custom operation.
- Bind group 0 is
SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR. - Bind group 1 is
SelectionBundle::BOX_BIND_GROUP_LAYOUT_DESCRIPTOR.
Trait Implementations§
Auto Trait Implementations§
impl<G> Freeze for SelectionBundle<G>
impl<G> !RefUnwindSafe for SelectionBundle<G>
impl<G> Send for SelectionBundle<G>
impl<G> Sync for SelectionBundle<G>
impl<G> Unpin for SelectionBundle<G>where
G: Unpin,
impl<G> !UnwindSafe for SelectionBundle<G>
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self>
fn into_either(self, into_left: bool) -> Either<Self, Self>
self into a Left variant of Either<Self, Self>
if into_left is true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
self into a Left variant of Either<Self, Self>
if into_left(&self) returns true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read more