Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: aabb/obb compare #36

Merged
merged 5 commits into from
Nov 26, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,11 @@ path = "tools/ply_to_gcloud.rs"
required-features = ["io_ply"]


[[bin]]
name = "compare_aabb_obb"
path = "tools/compare_aabb_obb.rs"


[[bench]]
name = "io"
harness = false
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ fn setup_gaussian_cloud(

- [ply to gcloud converter](tools/README.md#ply-to-gcloud-converter)
- [ ] gaussian cloud training tool
- aabb vs. obb gaussian comparison via `cargo run --bin compare_aabb_obb`

## wasm support

Expand Down
6 changes: 3 additions & 3 deletions src/gaussian.rs
Original file line number Diff line number Diff line change
Expand Up @@ -101,8 +101,6 @@ where
}


pub const MAX_SIZE_VARIANCE: f32 = 5.0;

#[derive(
Clone,
Debug,
Expand Down Expand Up @@ -195,6 +193,8 @@ impl GaussianCloud {
}
}

cloud.0.push(cloud.0[0]);

cloud
}
}
Expand All @@ -213,7 +213,7 @@ impl Default for GaussianCloudSettings {
fn default() -> Self {
Self {
aabb: false,
global_scale: 1.0,
global_scale: 2.0,
global_transform: Transform::IDENTITY.into(),
visualize_bounding_box: false,
}
Expand Down
3 changes: 2 additions & 1 deletion src/io/ply.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,11 +11,12 @@ use ply_rs::{
use crate::gaussian::{
Gaussian,
MAX_SH_COEFF_COUNT_PER_CHANNEL,
MAX_SIZE_VARIANCE,
SH_CHANNELS,
};


pub const MAX_SIZE_VARIANCE: f32 = 5.0;

impl PropertyAccess for Gaussian {
fn new() -> Self {
Gaussian::default()
Expand Down
1 change: 1 addition & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ pub mod utils;
pub struct GaussianSplattingBundle {
pub settings: GaussianCloudSettings,
pub cloud: Handle<GaussianCloud>,
pub visibility: Visibility,
}


Expand Down
105 changes: 49 additions & 56 deletions src/render/gaussian.wgsl
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,11 @@ fn compute_cov3d(scale: vec3<f32>, rotation: vec4<f32>) -> array<f32, 6> {
);
}

fn compute_cov2d(position: vec3<f32>, scale: vec3<f32>, rotation: vec4<f32>) -> vec3<f32> {
fn compute_cov2d(
position: vec3<f32>,
scale: vec3<f32>,
rotation: vec4<f32>
) -> vec3<f32> {
let cov3d = compute_cov3d(scale, rotation);
let Vrk = mat3x3(
cov3d[0], cov3d[1], cov3d[2],
Expand All @@ -69,30 +73,16 @@ fn compute_cov2d(position: vec3<f32>, scale: vec3<f32>, rotation: vec4<f32>) ->

var t = view.inverse_view * vec4<f32>(position, 1.0);

let focal_x = 600.0;
let focal_y = 600.0;

let fovy = 2.0 * atan(1.0 / view.projection[1][1]);
let fovx = 2.0 * atan(1.0 / view.projection[0][0]);
let tan_fovy = tan(fovy * 0.5);
let tan_fovx = tan(fovx * 0.5);

let limx = 1.3 * tan_fovx;
let limy = 1.3 * tan_fovy;
let txtz = t.x / t.z;
let tytz = t.y / t.z;
t.x = min(limx, max(-limx, txtz)) * t.z;
t.y = min(limy, max(-limy, tytz)) * t.z;
let device_pixel_ratio = 1.0;
let focal = vec2<f32>(
view.projection.x.x * device_pixel_ratio * view.viewport.z * 0.45,
view.projection.y.y * device_pixel_ratio * view.viewport.w * 0.45,
);

let s = 1.0 / (t.z * t.z);
let J = mat3x3(
focal_x / t.z,
0.0,
-(focal_x * t.x) / (t.z * t.z),

0.0,
-focal_y / t.z,
(focal_y * t.y) / (t.z * t.z),

focal.x / t.z, 0.0, -(focal.x * t.x) * s,
0.0, -focal.y / t.z, (focal.y * t.y) * s,
0.0, 0.0, 0.0,
);

Expand All @@ -113,26 +103,22 @@ fn compute_cov2d(position: vec3<f32>, scale: vec3<f32>, rotation: vec4<f32>) ->
return vec3<f32>(cov[0][0], cov[0][1], cov[1][1]);
}


fn get_bounding_box(
cov2d: vec3<f32>,
direction: vec2<f32>,
) -> vec4<f32> {
// return vec4<f32>(offset, uv);

// let det = cov2d.x * cov2d.z - cov2d.y * cov2d.y;
// let mid = 0.5 * (cov2d.x + cov2d.z);
// let lambda1 = mid + sqrt(max(0.1, mid * mid - det));
// let lambda2 = mid - sqrt(max(0.1, mid * mid - det));
// let x_axis_length = sqrt(lambda1);
// let y_axis_length = sqrt(lambda2);

let det = cov2d.x * cov2d.z - cov2d.y * cov2d.y;
let mid = 0.5 * (cov2d.x + cov2d.z);
var discriminant = max(0.0, mid * mid - det);
let trace = cov2d.x + cov2d.z;
let mid = 0.5 * trace;
let discriminant = max(0.0, mid * mid - det);

let term = sqrt(discriminant);

let lambda1 = mid + term;
let lambda2 = max(mid - term, 0.0);

let lambda1 = mid + sqrt(discriminant);
let lambda2 = mid - sqrt(discriminant);
let x_axis_length = sqrt(lambda1);
let y_axis_length = sqrt(lambda2);

Expand All @@ -141,52 +127,62 @@ fn get_bounding_box(
// creates a square AABB (inefficient fragment usage)
let radius_px = 3.5 * max(x_axis_length, y_axis_length);
let radius_ndc = vec2<f32>(
radius_px / view.viewport.z,
radius_px / view.viewport.w,
radius_px / view.viewport.zw,
);

return vec4<f32>(
2.0 * radius_ndc * direction,
radius_ndc * direction,
radius_px * direction,
);
#endif

#ifdef USE_OBB

let a = (cov2d.x - cov2d.z) * (cov2d.x - cov2d.z);
let b = sqrt(a + 4.0 * cov2d.y * cov2d.y);
let major_radius = sqrt((cov2d.x + cov2d.z + b) * 0.5);
let minor_radius = sqrt((cov2d.x + cov2d.z - b) * 0.5);

let bounds = 3.5 * vec2<f32>(
x_axis_length,
y_axis_length,
major_radius,
minor_radius,
);

// bounding box is aligned to the eigenvectors with proper width/height
// collapse unstable eigenvectors to circle
let threshold = 0.1;
if (abs(lambda1 - lambda2) < threshold) {
let circle = direction * max(x_axis_length, y_axis_length);
let circle = direction * max(bounds.x, bounds.y);
return vec4<f32>(
circle / view.viewport.zw,
circle
circle,
);
}


let eigvec1 = normalize(vec2<f32>(
cov2d.y,
lambda1 - cov2d.x
-cov2d.y,
lambda1 - cov2d.x,
));
let eigvec2 = vec2<f32>(
-eigvec1.y,
eigvec1.x
eigvec1.y,
-eigvec1.x
);

let rotation_matrix = mat2x2(
eigvec1.x, eigvec2.x,
eigvec1.y, eigvec2.y
let rotation_matrix = transpose(
mat2x2(
eigvec1,
eigvec2,
)
);

let scaled_vertex = direction * bounds;
let rotated_vertex = scaled_vertex * rotation_matrix;

let scaling_factor = 1.0 / view.viewport.zw;
let ndc_vertex = rotated_vertex * scaling_factor;

return vec4<f32>(
scaled_vertex * rotation_matrix / view.viewport.zw,
0.0, 0.0,
ndc_vertex,
rotated_vertex,
);
#endif
}
Expand Down Expand Up @@ -234,7 +230,6 @@ fn vs_points(

let cov2d = compute_cov2d(transformed_position, point.scale_opacity.rgb, point.rotation);

// TODO: remove conic when OBB is used
let det = cov2d.x * cov2d.z - cov2d.y * cov2d.y;
let det_inv = 1.0 / det;
let conic = vec3<f32>(
Expand All @@ -261,8 +256,6 @@ fn vs_points(

@fragment
fn fs_main(input: GaussianOutput) -> @location(0) vec4<f32> {
// TODO: draw gaussian without conic (OBB)

#ifdef USE_AABB
let d = -input.major_minor;
let conic = input.conic;
Expand Down
17 changes: 13 additions & 4 deletions src/render/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -766,6 +766,7 @@ pub fn extract_gaussians(
Query<(
Entity,
// &ComputedVisibility,
&Visibility,
&Handle<GaussianCloud>,
&GaussianCloudSettings,
)>,
Expand All @@ -774,7 +775,16 @@ pub fn extract_gaussians(
let mut commands_list = Vec::with_capacity(*prev_commands_len);
// let visible_gaussians = gaussians_query.iter().filter(|(_, vis, ..)| vis.is_visible());

for (entity, verticies, settings) in gaussians_query.iter() {
for (
entity,
visibility,
verticies,
settings,
) in gaussians_query.iter() {
if visibility == Visibility::Hidden {
continue;
}

let settings_uniform = GaussianCloudUniform {
transform: settings.global_transform.compute_matrix(),
global_scale: settings.global_scale,
Expand Down Expand Up @@ -805,7 +815,7 @@ pub struct GaussianCloudBindGroup {
pub sorted_bind_group: BindGroup,
}

pub fn queue_gaussian_bind_group(
fn queue_gaussian_bind_group(
mut commands: Commands,
mut groups: ResMut<GaussianUniformBindGroups>,
gaussian_cloud_pipeline: Res<GaussianCloudPipeline>,
Expand All @@ -831,7 +841,7 @@ pub fn queue_gaussian_bind_group(
resource: BindingResource::Buffer(BufferBinding {
buffer: model,
offset: 0,
size: BufferSize::new(model.size()),
size: GaussianCloudUniform::min_size().into(),
}),
},
],
Expand Down Expand Up @@ -958,7 +968,6 @@ pub fn queue_gaussian_bind_group(
}
}


#[derive(Component)]
pub struct GaussianViewBindGroup {
pub value: BindGroup,
Expand Down
Loading