feat: Enable generic attributes for point cloud encode/decode (#61)

* Ensure attribute `unique_id`s are always unique

* Enable generic attributes for point clouds in addition to meshes

* Check that all `unique_ids` are >= -1

* Raise errors instead of silently failing on invalid input for encoding

Author: schmrlng

src/DracoPy.h

@@ -5,6 +5,7 @@
 #include<cmath>
 #include<vector>
 #include<cstddef>
+#include<sstream>
 #include "draco/compression/decode.h"
 #include "draco/compression/encode.h"
 #include "draco/compression/config/compression_shared.h"
@@ -291,6 +292,17 @@ namespace DracoFunctions {
     //    manually.
     draco::Mesh mesh; //Initialize a draco mesh
 
+    // Ensure unique ids for native attributes don't conflict with generic attributes
+    uint32_t next_unique_id = 0;
+    auto min_unique_id = std::min_element(unique_ids.begin(), unique_ids.end());
+    if (min_unique_id != unique_ids.end() && *min_unique_id < -1) {
+      throw std::invalid_argument("Should never be reached; all unique_ids should be >= -1.");
+    }
+    auto max_unique_id = std::max_element(unique_ids.begin(), unique_ids.end());
+    if (max_unique_id != unique_ids.end()) {
+      next_unique_id = *max_unique_id + 1;
+    }
+
     // Process vertices
     const size_t num_pts = points.size() / 3;
     mesh.set_num_points(num_pts);
@@ -331,6 +343,8 @@ namespace DracoFunctions {
                        sizeof(uint8_t) * colors_channel,   // byte stride
                        0);                                 // byte offset
       color_att_id = mesh.AddAttribute(colors_attr, true, num_pts);
+      mesh.attribute(color_att_id)->set_unique_id(next_unique_id);
+      next_unique_id++;
     }
     int tex_coord_att_id = -1;
     if(tex_coord_channel) {
@@ -343,6 +357,8 @@ namespace DracoFunctions {
                           sizeof(float) * tex_coord_channel,   // byte stride
                           0);                                  // byte offset
       tex_coord_att_id = mesh.AddAttribute(tex_coord_attr, true, num_pts);
+      mesh.attribute(tex_coord_att_id)->set_unique_id(next_unique_id);
+      next_unique_id++;
     }
 
     int normal_att_id = -1;
@@ -356,6 +372,8 @@ namespace DracoFunctions {
                        sizeof(float) * 3,                   // byte stride
                        0);                                  // byte offset
       normal_att_id = mesh.AddAttribute(normal_attr, true, num_pts);
+      mesh.attribute(normal_att_id)->set_unique_id(next_unique_id);
+      next_unique_id++;
     }
 
 
@@ -369,21 +387,22 @@ namespace DracoFunctions {
       draco::DataType dtype = static_cast<draco::DataType>(attr_data_types[i]);
       int num_components = attr_num_components[i];
       if (dtype != draco::DT_FLOAT32 && dtype != draco::DT_UINT8 && dtype != draco::DT_UINT16 && dtype != draco::DT_UINT32) {
-        // Unsupported data type, skip
-        // std::cout << "DEBUG: Unsupported attribute data type for " << unique_ids[i] << std::endl;
-        generic_attr_ids.push_back(-1);
-        continue;
+        std::ostringstream oss;
+        oss << "Unsupported attribute data type for attribute with unique id " << int(unique_ids[i]);
+        throw std::invalid_argument(oss.str());
       }
       generic_attr.Init(draco::GeometryAttribute::GENERIC, nullptr, num_components, dtype, false, 0, 0);
       int att_id = mesh.AddAttribute(generic_attr, true, num_pts);
       if (att_id == -1) {
-        // Failed to add attribute, skip
-        // std::cout << "DEBUG: Failed to add attribute " << unique_ids[i] << std::endl;
-        generic_attr_ids.push_back(-1);
-        continue;
+        std::ostringstream oss;
+        oss << "Failed to add attribute with unique id " << int(unique_ids[i]);
+        throw std::runtime_error(oss.str());
       }
       if (unique_ids[i] >= 0) {
         mesh.attribute(att_id)->set_unique_id(unique_ids[i]);
+      } else {
+        mesh.attribute(att_id)->set_unique_id(next_unique_id);
+        next_unique_id++;
       }
       if (!attr_names[i].empty()) {
         auto attribute_metadata = std::unique_ptr<draco::AttributeMetadata>(new draco::AttributeMetadata());
@@ -396,6 +415,8 @@ namespace DracoFunctions {
 
 
     const int pos_att_id = mesh.AddAttribute(positions_attr, true, num_pts);
+    mesh.attribute(pos_att_id)->set_unique_id(next_unique_id);
+    next_unique_id++;
     std::vector<int32_t> pts_int32;
     std::vector<uint32_t> pts_uint32;
     if (integer_mark == 1) {
@@ -439,8 +460,7 @@ namespace DracoFunctions {
         const auto &uint32_data = attr_uint32_data[j];
 
         if (generic_attr_ids[j] == -1) {
-          // Skip if attribute was not added
-          continue;
+          throw std::runtime_error("Should never be reached; all successfully added attributes should have nonnegative att_id.");
         } 
 
         if (attr_data_types[j] == draco::DT_FLOAT32) {
@@ -500,12 +520,31 @@ namespace DracoFunctions {
     const float *quantization_origin, const bool preserve_order,
     const bool create_metadata, const int integer_mark,
     const std::vector<uint8_t> &colors,
-    const uint8_t colors_channel
+    const uint8_t colors_channel,
+    std::vector<int8_t>& unique_ids,
+    std::vector<std::vector<float>>& attr_float_data,
+    std::vector<std::vector<uint8_t>>& attr_uint8_data,
+    std::vector<std::vector<uint16_t>>& attr_uint16_data,
+    std::vector<std::vector<uint32_t>>& attr_uint32_data,
+    std::vector<int>& attr_data_types,
+    std::vector<int>& attr_num_components,
+    std::vector<std::string>& attr_names
   ) {
     int num_points = points.size() / 3;
     draco::PointCloudBuilder pcb;
     pcb.Start(num_points);
 
+    // Ensure unique ids for native attributes don't conflict with generic attributes
+    uint32_t next_unique_id = 0;
+    auto min_unique_id = std::min_element(unique_ids.begin(), unique_ids.end());
+    if (min_unique_id != unique_ids.end() && *min_unique_id < -1) {
+      throw std::invalid_argument("Should never be reached; all unique_ids should be >= -1.");
+    }
+    auto max_unique_id = std::max_element(unique_ids.begin(), unique_ids.end());
+    if (max_unique_id != unique_ids.end()) {
+      next_unique_id = *max_unique_id + 1;
+    }
+
     auto dtype = (integer_mark == 1)
       ? draco::DataType::DT_INT32
       : (
@@ -517,11 +556,15 @@ namespace DracoFunctions {
     const int pos_att_id = pcb.AddAttribute(
       draco::GeometryAttribute::POSITION, 3, dtype
     );
+    pcb.SetAttributeUniqueId(pos_att_id, next_unique_id);
+    next_unique_id++;
 
     if(colors_channel){
       const int color_att_id = pcb.AddAttribute(
         draco::GeometryAttribute::COLOR, colors_channel, draco::DataType::DT_UINT8
       );
+      pcb.SetAttributeUniqueId(color_att_id, next_unique_id);
+      next_unique_id++;
       for (draco::PointIndex i(0); i < num_points; i++) {
         pcb.SetAttributeValueForPoint(pos_att_id, i, points.data() + 3 * i.value());
         pcb.SetAttributeValueForPoint(color_att_id, i, colors.data() + colors_channel * i.value());
@@ -532,6 +575,44 @@ namespace DracoFunctions {
       }
     }
 
+    // GENERIC ATTRIBUTES
+    for (size_t j = 0; j < unique_ids.size(); ++j) {
+      draco::DataType dtype = static_cast<draco::DataType>(attr_data_types[j]);
+      int num_components = attr_num_components[j];
+      if (dtype != draco::DT_FLOAT32 && dtype != draco::DT_UINT8 && dtype != draco::DT_UINT16 && dtype != draco::DT_UINT32) {
+        std::ostringstream oss;
+        oss << "Unsupported attribute data type for attribute with unique id " << int(unique_ids[j]);
+        throw std::invalid_argument(oss.str());
+      }
+      int att_id = pcb.AddAttribute(draco::GeometryAttribute::GENERIC, num_components, dtype);
+      if (att_id == -1) {
+        std::ostringstream oss;
+        oss << "Failed to add attribute with unique id " << int(unique_ids[j]);
+        throw std::runtime_error(oss.str());
+      }
+      if (unique_ids[j] >= 0) {
+        pcb.SetAttributeUniqueId(att_id, unique_ids[j]);
+      } else {
+        pcb.SetAttributeUniqueId(att_id, next_unique_id);
+        next_unique_id++;
+      }
+      if (!attr_names[j].empty()) {
+        auto attribute_metadata = std::unique_ptr<draco::AttributeMetadata>(new draco::AttributeMetadata());
+        attribute_metadata->AddEntryString("name", attr_names[j]);
+        pcb.AddAttributeMetadata(att_id, std::move(attribute_metadata));
+      }
+
+      if (dtype == draco::DT_FLOAT32) {
+        pcb.SetAttributeValuesForAllPoints(att_id, attr_float_data[j].data(), 0);
+      } else if (dtype == draco::DT_UINT8) {
+        pcb.SetAttributeValuesForAllPoints(att_id, attr_uint8_data[j].data(), 0);
+      } else if (dtype == draco::DT_UINT16) {
+        pcb.SetAttributeValuesForAllPoints(att_id, attr_uint16_data[j].data(), 0);
+      } else if (dtype == draco::DT_UINT32) {
+        pcb.SetAttributeValuesForAllPoints(att_id, attr_uint32_data[j].data(), 0);
+      }
+    }
+
     std::unique_ptr<draco::PointCloud> ptr_point_cloud = pcb.Finalize(!preserve_order);
     draco::PointCloud *point_cloud = ptr_point_cloud.get();
     draco::Encoder encoder;

src/DracoPy.pxd

@@ -95,5 +95,13 @@ cdef extern from "DracoPy.h" namespace "DracoFunctions":
         const bool create_metadata,
         const int integer_mark,
         const vector[uint8_t] colors,
-        const uint8_t colors_channel
+        const uint8_t colors_channel,
+        vector[int8_t]& unique_ids,
+        vector[vector[float]]& attr_float_data,
+        vector[vector[uint8_t]]& attr_uint8_data,
+        vector[vector[uint16_t]]& attr_uint16_data,
+        vector[vector[uint32_t]]& attr_uint32_data,
+        vector[int]& attr_data_types,
+        vector[int]& attr_num_components,
+        vector[string]& attr_names
     ) except +

src/DracoPy.pyx

@@ -405,7 +405,11 @@ def encode(
             pointsview, quantization_bits, compression_level,
             quantization_range, <float*>&quant_origin[0],
             preserve_order, create_metadata, integer_mark,
-            colorsview, colors_channel
+            colorsview, colors_channel,
+            unique_ids, attr_float_data, attr_uint8_data,
+            attr_uint16_data, attr_uint32_data,
+            attr_data_types, attr_num_components,
+            attr_names
         )
     else:
         facesview = faces.reshape((faces.size,))

tests.py

@@ -268,7 +268,8 @@ def test_normals_encoding():
     assert np.allclose(decoded_mesh.normals, test_normals)
 
 
-def test_generic_attributes():
+@pytest.mark.parametrize("object_type", [DracoPy.DracoMesh, DracoPy.DracoPointCloud])
+def test_generic_attributes(object_type):
     # Read reference mesh
     with open(os.path.join(testdata_directory, "bunny.drc"), 'rb') as draco_file:
         mesh = DracoPy.decode(draco_file.read())
@@ -282,16 +283,26 @@ def test_generic_attributes():
     generic_attributes = {
         0: test_tangents,
         1: test_joints,
-        6: test_weights
+        3: test_weights
     }
-    binary = DracoPy.encode(mesh.points, mesh.faces, generic_attributes=generic_attributes)
+    if object_type is DracoPy.DracoMesh:
+        binary = DracoPy.encode(mesh.points, mesh.faces, generic_attributes=generic_attributes)
+    else:
+        binary = DracoPy.encode(mesh.points, generic_attributes=generic_attributes)
 
     # Decode and verify attributes
-    decoded_mesh = DracoPy.decode(binary)
-
-    decoded_tangents = decoded_mesh.get_attribute_by_unique_id(0)
-    decoded_joints = decoded_mesh.get_attribute_by_unique_id(1)
-    decoded_weights = decoded_mesh.get_attribute_by_unique_id(6)
+    decoded_object = DracoPy.decode(binary)
+    if object_type is DracoPy.DracoMesh:
+        assert type(decoded_object) is DracoPy.DracoMesh
+    else:
+        assert type(decoded_object) is DracoPy.DracoPointCloud
+
+    assert len(decoded_object.attributes) == len(
+        set([attr["unique_id"] for attr in decoded_object.attributes])
+    )
+    decoded_tangents = decoded_object.get_attribute_by_unique_id(0)
+    decoded_joints = decoded_object.get_attribute_by_unique_id(1)
+    decoded_weights = decoded_object.get_attribute_by_unique_id(3)
 
     assert decoded_tangents["attribute_type"] == DracoPy.AttributeType.GENERIC
     assert decoded_tangents["data_type"] == DracoPy.DataType.DT_FLOAT32
@@ -309,7 +320,8 @@ def test_generic_attributes():
     assert np.allclose(decoded_weights["data"], test_weights)
 
 
-def test_named_generic_attributes():
+@pytest.mark.parametrize("object_type", [DracoPy.DracoMesh, DracoPy.DracoPointCloud])
+def test_named_generic_attributes(object_type):
     # Read reference mesh
     with open(os.path.join(testdata_directory, "bunny.drc"), 'rb') as draco_file:
         mesh = DracoPy.decode(draco_file.read())
@@ -325,14 +337,24 @@ def test_named_generic_attributes():
         'joints': test_joints,
         'weights': test_weights
     }
-    binary = DracoPy.encode(mesh.points, mesh.faces, generic_attributes=generic_attributes)
+    if object_type is DracoPy.DracoMesh:
+        binary = DracoPy.encode(mesh.points, mesh.faces, generic_attributes=generic_attributes)
+    else:
+        binary = DracoPy.encode(mesh.points, generic_attributes=generic_attributes)
 
     # Decode and verify attributes
-    decoded_mesh = DracoPy.decode(binary)
-
-    decoded_tangents = decoded_mesh.get_attribute_by_name('tangents')
-    decoded_joints = decoded_mesh.get_attribute_by_name('joints')
-    decoded_weights = decoded_mesh.get_attribute_by_name('weights')
+    decoded_object = DracoPy.decode(binary)
+    if object_type is DracoPy.DracoMesh:
+        assert type(decoded_object) is DracoPy.DracoMesh
+    else:
+        assert type(decoded_object) is DracoPy.DracoPointCloud
+    
+    assert len(decoded_object.attributes) == len(
+        set([attr["unique_id"] for attr in decoded_object.attributes])
+    )
+    decoded_tangents = decoded_object.get_attribute_by_name('tangents')
+    decoded_joints = decoded_object.get_attribute_by_name('joints')
+    decoded_weights = decoded_object.get_attribute_by_name('weights')
 
     assert decoded_tangents["attribute_type"] == DracoPy.AttributeType.GENERIC
     assert decoded_tangents["data_type"] == DracoPy.DataType.DT_FLOAT32