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

Add support for mesh vertex colors #1671

Merged
merged 26 commits into from
Mar 23, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
26 commits
Select commit Hold shift + click to select a range
df1b60b
Refactor: use `izip!`
emilk Mar 22, 2023
91753d6
Improve docstrings for rr.log_mesh and rr.log_mesh_file
emilk Mar 22, 2023
e6487a9
Add support for mesh vertex colors in Python API
emilk Mar 22, 2023
21454e3
Improve sanity checking for raw meshes
emilk Mar 22, 2023
43e38fb
Add vertex colors to RawMesh3D
emilk Mar 22, 2023
332eaf0
Rename normals to vertex_normals
emilk Mar 22, 2023
d7f190e
Split vertex normals and texcoords into different vertex buffers
emilk Mar 22, 2023
3a75372
Add sanity check for shader locations
emilk Mar 22, 2023
2a3f031
Pipe vertex colors into the GPU
emilk Mar 22, 2023
25a9609
Use mesh vertex colors in shader
emilk Mar 22, 2023
283940e
Don't crash on invalid meshes
emilk Mar 22, 2023
a155b8a
Fix num_triangles function
emilk Mar 22, 2023
dbb4d80
Better sanity checking of loaded GLB/GLTF and OBJ meshes
emilk Mar 22, 2023
087037b
Cleanup
emilk Mar 22, 2023
c8e352d
Fix typo
emilk Mar 22, 2023
40b219e
Handle RGB colors too
emilk Mar 22, 2023
956f6dd
Fix broken test
emilk Mar 22, 2023
2ff4963
Better names
emilk Mar 22, 2023
4ae1240
Fix mypy types
emilk Mar 22, 2023
83c82ae
Document the color type of point cloud shaders
emilk Mar 23, 2023
ab6cd29
fix typo
emilk Mar 23, 2023
8c12b53
document additive_tint_rgb
emilk Mar 23, 2023
447b9b3
fix TODO
emilk Mar 23, 2023
347f40a
Add Rgba32Unmul color type
emilk Mar 23, 2023
9f21f2d
Create a helper function for generating the VertexBufferLayouts
emilk Mar 23, 2023
71b9154
Use from_rgb helper
emilk Mar 23, 2023
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
21 changes: 16 additions & 5 deletions crates/re_log_types/src/component_types/color.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
use arrow2_convert::{ArrowDeserialize, ArrowField, ArrowSerialize};

use crate::msg_bundle::Component;

/// An RGBA color tuple.
/// An RGBA color tuple with unmultiplied/separate alpha,
/// in sRGB gamma space with linear alpha.
///
/// ```
/// use re_log_types::component_types::ColorRGBA;
Expand All @@ -11,9 +10,21 @@ use crate::msg_bundle::Component;
///
/// assert_eq!(ColorRGBA::data_type(), DataType::UInt32);
/// ```
#[derive(Clone, Copy, Debug, PartialEq, Eq, ArrowField, ArrowSerialize, ArrowDeserialize)]
#[derive(
Clone,
Copy,
Debug,
PartialEq,
Eq,
bytemuck::Pod,
bytemuck::Zeroable,
arrow2_convert::ArrowField,
arrow2_convert::ArrowSerialize,
arrow2_convert::ArrowDeserialize,
)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
#[arrow_field(transparent)]
#[repr(transparent)]
pub struct ColorRGBA(pub u32);

impl ColorRGBA {
Expand All @@ -23,7 +34,7 @@ impl ColorRGBA {
}

#[inline]
pub fn from_rgba(r: u8, g: u8, b: u8, a: u8) -> Self {
pub fn from_unmultiplied_rgba(r: u8, g: u8, b: u8, a: u8) -> Self {
Self::from([r, g, b, a])
}

Expand Down
163 changes: 108 additions & 55 deletions crates/re_log_types/src/component_types/mesh3d.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use arrow2_convert::{serialize::ArrowSerialize, ArrowDeserialize, ArrowField, Ar

use crate::msg_bundle::Component;

use super::{FieldError, Vec4D};
use super::{ColorRGBA, FieldError, Vec4D};

// ----------------------------------------------------------------------------

Expand Down Expand Up @@ -88,12 +88,18 @@ impl ArrowDeserialize for MeshId {

#[derive(thiserror::Error, Debug)]
pub enum RawMeshError {
#[error("Positions array length must be divisible by 3 (triangle list), got {0}")]
#[error("Positions array length must be divisible by 3 (xyz, xyz, …), got {0}")]
PositionsNotDivisibleBy3(usize),

#[error("Indices array length must be divisible by 3 (triangle list), got {0}")]
IndicesNotDivisibleBy3(usize),

#[error("No indices were specified, so the number of positions must be divisible by 9 [(xyz xyz xyz), …], got {0}")]
PositionsAreNotTriangles(usize),

#[error("Index out of bounds: got index={index} with {num_vertices} vertices")]
IndexOutOfBounds { index: u32, num_vertices: usize },

#[error(
"Positions & normals array must have the same length, \
got positions={0} vs. normals={1}"
Expand All @@ -111,15 +117,18 @@ pub enum RawMeshError {
/// RawMesh3D::data_type(),
/// DataType::Struct(vec![
/// Field::new("mesh_id", DataType::FixedSizeBinary(16), false),
/// Field::new("positions", DataType::List(Box::new(
/// Field::new("vertex_positions", DataType::List(Box::new(
/// Field::new("item", DataType::Float32, false)),
/// ), false),
/// Field::new("indices", DataType::List(Box::new(
/// Field::new("vertex_colors", DataType::List(Box::new(
/// Field::new("item", DataType::UInt32, false)),
/// ), true),
/// Field::new("normals", DataType::List(Box::new(
/// Field::new("vertex_normals", DataType::List(Box::new(
/// Field::new("item", DataType::Float32, false)),
/// ), true),
/// Field::new("indices", DataType::List(Box::new(
/// Field::new("item", DataType::UInt32, false)),
/// ), true),
/// Field::new("albedo_factor", DataType::FixedSizeList(
/// Box::new(Field::new("item", DataType::Float32, false)),
/// 4
Expand All @@ -132,22 +141,27 @@ pub enum RawMeshError {
pub struct RawMesh3D {
pub mesh_id: MeshId,

/// The flattened positions array of this mesh.
/// The flattened vertex positions array of this mesh.
///
/// Meshes are always triangle lists, i.e. the length of this vector should always be
/// divisible by 3.
pub positions: Vec<f32>,

/// Optionally, the flattened indices array for this mesh.
/// The length of this vector should always be divisible by three (since this is a 3D mesh).
///
/// Meshes are always triangle lists, i.e. the length of this vector should always be
/// divisible by 3.
pub indices: Option<Vec<u32>>,
/// If no indices are specified, then each triplet of vertex positions are intrpreted as a triangle
/// and the length of this must be divisible by 9.
pub vertex_positions: Vec<f32>,

/// Per-vertex albedo colors.
pub vertex_colors: Option<Vec<ColorRGBA>>,

/// Optionally, the flattened normals array for this mesh.
///
/// If specified, this must match the length of `Self::positions`.
pub normals: Option<Vec<f32>>,
pub vertex_normals: Option<Vec<f32>>,

/// Optionally, the flattened indices array for this mesh.
///
/// Meshes are always triangle lists, i.e. the length of this vector should always be
/// divisible by three.
pub indices: Option<Vec<u32>>,

/// Albedo factor applied to the final color of the mesh.
///
Expand All @@ -163,20 +177,37 @@ pub struct RawMesh3D {

impl RawMesh3D {
pub fn sanity_check(&self) -> Result<(), RawMeshError> {
if self.positions.len() % 3 != 0 {
return Err(RawMeshError::PositionsNotDivisibleBy3(self.positions.len()));
if self.vertex_positions.len() % 3 != 0 {
return Err(RawMeshError::PositionsNotDivisibleBy3(
self.vertex_positions.len(),
));
}

let num_vertices = self.vertex_positions.len() / 3;

if let Some(indices) = &self.indices {
if indices.len() % 3 != 0 {
return Err(RawMeshError::IndicesNotDivisibleBy3(indices.len()));
}

for &index in indices {
if num_vertices <= index as usize {
return Err(RawMeshError::IndexOutOfBounds {
index,
num_vertices,
});
}
}
} else if self.vertex_positions.len() % 9 != 0 {
return Err(RawMeshError::PositionsAreNotTriangles(
self.vertex_positions.len(),
));
}

if let Some(normals) = &self.normals {
if normals.len() != self.positions.len() {
if let Some(normals) = &self.vertex_normals {
if normals.len() != self.vertex_positions.len() {
return Err(RawMeshError::MismatchedPositionsNormals(
self.positions.len(),
self.vertex_positions.len(),
normals.len(),
));
}
Expand All @@ -186,11 +217,17 @@ impl RawMesh3D {
}

#[inline]
pub fn num_triangles(&self) -> usize {
#[cfg(debug_assertions)]
self.sanity_check().unwrap();
pub fn num_vertices(&self) -> usize {
self.vertex_positions.len() / 3
}

self.positions.len() / 3
#[inline]
pub fn num_triangles(&self) -> usize {
if let Some(indices) = &self.indices {
indices.len() / 3
} else {
self.num_vertices() / 3
}
}
}

Expand Down Expand Up @@ -392,40 +429,56 @@ impl Mesh3D {
}
}

#[test]
fn test_mesh_roundtrip() {
use arrow2::array::Array;
use arrow2_convert::{deserialize::TryIntoCollection, serialize::TryIntoArrow};
#[cfg(test)]
mod tests {
use super::*;

// Encoded
{
let mesh_in = vec![Mesh3D::Encoded(EncodedMesh3D {
fn example_raw_mesh() -> RawMesh3D {
let mesh = RawMesh3D {
mesh_id: MeshId::random(),
format: MeshFormat::Glb,
bytes: std::sync::Arc::new([5, 9, 13, 95, 38, 42, 98, 17]),
transform: [
[1.0, 2.0, 3.0],
[4.0, 5.0, 6.0],
[7.0, 8.0, 9.0],
[10.0, 11.0, 12.],
],
})];
let array: Box<dyn Array> = mesh_in.try_into_arrow().unwrap();
let mesh_out: Vec<Mesh3D> = TryIntoCollection::try_into_collection(array).unwrap();
assert_eq!(mesh_in, mesh_out);
vertex_positions: vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 8.0, 9.0, 10.0],
vertex_colors: Some(vec![
ColorRGBA(0xff0000ff),
ColorRGBA(0x00ff00ff),
ColorRGBA(0x0000ffff),
]),
indices: vec![0, 1, 2].into(),
vertex_normals: vec![10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 80.0, 90.0, 100.0].into(),
albedo_factor: Vec4D([0.5, 0.5, 0.5, 1.0]).into(),
};
mesh.sanity_check().unwrap();
mesh
}

// Raw
{
let mesh_in = vec![Mesh3D::Raw(RawMesh3D {
mesh_id: MeshId::random(),
positions: vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 8.0, 9.0, 10.0],
indices: vec![1, 2, 3].into(),
normals: vec![10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 80.0, 90.0, 100.0].into(),
albedo_factor: Vec4D([0.5, 0.5, 0.5, 1.0]).into(),
})];
let array: Box<dyn Array> = mesh_in.try_into_arrow().unwrap();
let mesh_out: Vec<Mesh3D> = TryIntoCollection::try_into_collection(array).unwrap();
assert_eq!(mesh_in, mesh_out);
#[test]
fn test_mesh_roundtrip() {
use arrow2::array::Array;
use arrow2_convert::{deserialize::TryIntoCollection, serialize::TryIntoArrow};

// Encoded
{
let mesh_in = vec![Mesh3D::Encoded(EncodedMesh3D {
mesh_id: MeshId::random(),
format: MeshFormat::Glb,
bytes: std::sync::Arc::new([5, 9, 13, 95, 38, 42, 98, 17]),
transform: [
[1.0, 2.0, 3.0],
[4.0, 5.0, 6.0],
[7.0, 8.0, 9.0],
[10.0, 11.0, 12.],
],
})];
let array: Box<dyn Array> = mesh_in.try_into_arrow().unwrap();
let mesh_out: Vec<Mesh3D> = TryIntoCollection::try_into_collection(array).unwrap();
assert_eq!(mesh_in, mesh_out);
}

// Raw
{
let mesh_in = vec![Mesh3D::Raw(example_raw_mesh())];
let array: Box<dyn Array> = mesh_in.try_into_arrow().unwrap();
let mesh_out: Vec<Mesh3D> = TryIntoCollection::try_into_collection(array).unwrap();
assert_eq!(mesh_in, mesh_out);
}
}
}
33 changes: 21 additions & 12 deletions crates/re_renderer/shader/instanced_mesh.wgsl
Original file line number Diff line number Diff line change
Expand Up @@ -10,16 +10,17 @@ var albedo_texture: texture_2d<f32>;
struct MaterialUniformBuffer {
albedo_factor: Vec4,
};

@group(1) @binding(1)
var<uniform> material: MaterialUniformBuffer;

struct VertexOut {
@builtin(position) position: Vec4,
@location(0) texcoord: Vec2,
@location(1) normal_world_space: Vec3,
@location(2) additive_tint_rgb: Vec3,

@location(3) @interpolate(flat)
@location(0) color: Vec4, // 0-1 linear space with unmultiplied/separate alpha
@location(1) texcoord: Vec2,
@location(2) normal_world_space: Vec3,
@location(3) additive_tint_rgb: Vec3, // 0-1 linear space
@location(4) @interpolate(flat)
outline_mask_ids: UVec2,
};

Expand All @@ -38,6 +39,7 @@ fn vs_main(in_vertex: VertexIn, in_instance: InstanceIn) -> VertexOut {

var out: VertexOut;
out.position = frame.projection_from_world * Vec4(world_position, 1.0);
out.color = linear_from_srgba(in_vertex.color);
out.texcoord = in_vertex.texcoord;
out.normal_world_space = world_normal;
out.additive_tint_rgb = linear_from_srgb(in_instance.additive_tint_srgb.rgb);
Expand All @@ -49,16 +51,23 @@ fn vs_main(in_vertex: VertexIn, in_instance: InstanceIn) -> VertexOut {
@fragment
fn fs_main_shaded(in: VertexOut) -> @location(0) Vec4 {
let albedo = textureSample(albedo_texture, trilinear_sampler, in.texcoord).rgb
* material.albedo_factor.rgb + in.additive_tint_rgb;
* in.color.rgb
* material.albedo_factor.rgb
+ in.additive_tint_rgb;

// Hardcoded lambert lighting. TODO(andreas): Some microfacet model.
let light_dir = normalize(vec3(1.0, 2.0, 0.0)); // TODO(andreas): proper lighting
let normal = normalize(in.normal_world_space);
let shading = clamp(dot(normal, light_dir), 0.0, 1.0) + 0.2;
if (all(in.normal_world_space == Vec3(0.0, 0.0, 0.0))) {
emilk marked this conversation as resolved.
Show resolved Hide resolved
// no normal, no shading
return Vec4(albedo, 1.0);
} else {
// Hardcoded lambert lighting. TODO(andreas): Some microfacet model.
let light_dir = normalize(vec3(1.0, 2.0, 0.0)); // TODO(andreas): proper lighting
let normal = normalize(in.normal_world_space);
let shading = clamp(dot(normal, light_dir), 0.0, 1.0) + 0.2;

let radiance = albedo * shading;
let radiance = albedo * shading;

return Vec4(radiance, 1.0);
return Vec4(radiance, 1.0);
}
}

@fragment
Expand Down
21 changes: 11 additions & 10 deletions crates/re_renderer/shader/mesh_vertex.wgsl
Original file line number Diff line number Diff line change
@@ -1,20 +1,21 @@
// See mesh.rs#MeshVertex
struct VertexIn {
@location(0) position: Vec3,
@location(1) normal: Vec3,
@location(2) texcoord: Vec2,
@location(1) color: Vec4, // gamma-space 0-1, unmultiplied
@location(2) normal: Vec3,
@location(3) texcoord: Vec2,
};

// See mesh_renderer.rs
struct InstanceIn {
// We could alternatively store projection_from_mesh, but world position might be useful
// in the future and this saves us a Vec4 and simplifies dataflow on the cpu side.
@location(3) world_from_mesh_row_0: Vec4,
@location(4) world_from_mesh_row_1: Vec4,
@location(5) world_from_mesh_row_2: Vec4,
@location(6) world_from_mesh_normal_row_0: Vec3,
@location(7) world_from_mesh_normal_row_1: Vec3,
@location(8) world_from_mesh_normal_row_2: Vec3,
@location(9) additive_tint_srgb: Vec4,
@location(10) outline_mask_ids: UVec2,
@location(4) world_from_mesh_row_0: Vec4,
@location(5) world_from_mesh_row_1: Vec4,
@location(6) world_from_mesh_row_2: Vec4,
@location(7) world_from_mesh_normal_row_0: Vec3,
@location(8) world_from_mesh_normal_row_1: Vec3,
@location(9) world_from_mesh_normal_row_2: Vec3,
@location(10) additive_tint_srgb: Vec4,
@location(11) outline_mask_ids: UVec2,
};
2 changes: 1 addition & 1 deletion crates/re_renderer/shader/point_cloud.wgsl
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ var<private> TEXTURE_SIZE: i32 = 2048;

struct VertexOut {
@builtin(position) position: Vec4,
@location(0) color: Vec4,
@location(0) color: Vec4, // linear RGBA with unmulitplied/separate alpha
@location(1) world_position: Vec3,
@location(2) point_center: Vec3,
@location(3) radius: f32,
Expand Down
Loading