Add support for mesh vertex colors (#1671)

* Refactor: use `izip!`

* Improve docstrings for rr.log_mesh and rr.log_mesh_file

* Add support for mesh vertex colors in Python API

* Improve sanity checking for raw meshes

* Add vertex colors to RawMesh3D

* Rename normals to vertex_normals

* Split vertex normals and texcoords into different vertex buffers

* Add sanity check for shader locations

* Pipe vertex colors into the GPU

* Use mesh vertex colors in shader

* Don't crash on invalid meshes

* Fix num_triangles function

* Better sanity checking of loaded GLB/GLTF and OBJ meshes

* Cleanup

* Fix typo

* Handle RGB colors too

* Fix broken test

* Better names

* Fix mypy types

* Document the color type of point cloud shaders

* fix typo

Co-authored-by: Andreas Reich <[email protected]>

* document additive_tint_rgb

Co-authored-by: Andreas Reich <[email protected]>

* fix TODO

Co-authored-by: Andreas Reich <[email protected]>

* Add Rgba32Unmul color type

* Create a helper function for generating the VertexBufferLayouts

* Use from_rgb helper

---------

Co-authored-by: Andreas Reich <[email protected]>

crates/re_log_types/src/component_types/color.rs

@@ -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;
@@ -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 {
@@ -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])
     }
 

crates/re_log_types/src/component_types/mesh3d.rs

@@ -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};
 
 // ----------------------------------------------------------------------------
 
@@ -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}"
@@ -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
@@ -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.
     ///
@@ -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(),
                 ));
             }
@@ -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
+        }
     }
 }
 
@@ -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);
+        }
     }
 }

crates/re_renderer/shader/instanced_mesh.wgsl

@@ -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,
 };
 
@@ -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);
@@ -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))) {
+        // 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

crates/re_renderer/shader/mesh_vertex.wgsl

@@ -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,
 };

crates/re_renderer/shader/point_cloud.wgsl

@@ -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,