[Bf-extensions-cvs] [4f8ccce4] master: glTF importer: Import custom primitive attributes as mesh attributes

[Julien Duroure]

Commit: 4f8ccce4a5357974804dcecce88deddd67dcb8f3
Author: Julien Duroure
Date:   Wed Jan 18 17:57:19 2023 +0100
Branches: master
https://developer.blender.org/rBA4f8ccce4a5357974804dcecce88deddd67dcb8f3

glTF importer: Import custom primitive attributes as mesh attributes

===================================================================

M	io_scene_gltf2/__init__.py
M	io_scene_gltf2/blender/com/gltf2_blender_conversion.py
M	io_scene_gltf2/blender/imp/gltf2_blender_mesh.py

===================================================================

diff --git a/io_scene_gltf2/__init__.py b/io_scene_gltf2/__init__.py
index 95f22981..dfd0e426 100755
--- a/io_scene_gltf2/__init__.py
+++ b/io_scene_gltf2/__init__.py
@@ -4,7 +4,7 @@
 bl_info = {
     'name': 'glTF 2.0 format',
     'author': 'Julien Duroure, Scurest, Norbert Nopper, Urs Hanselmann, Moritz Becher, Benjamin Schmithüsen, Jim Eckerlein, and many external contributors',
-    "version": (3, 5, 16),
+    "version": (3, 5, 17),
     'blender': (3, 4, 0),
     'location': 'File > Import-Export',
     'description': 'Import-Export as glTF 2.0',
diff --git a/io_scene_gltf2/blender/com/gltf2_blender_conversion.py b/io_scene_gltf2/blender/com/gltf2_blender_conversion.py
index 8b1e5452..4cc46e5c 100755
--- a/io_scene_gltf2/blender/com/gltf2_blender_conversion.py
+++ b/io_scene_gltf2/blender/com/gltf2_blender_conversion.py
@@ -105,3 +105,24 @@ def get_numpy_type(attribute_component_type):
         "FLOAT": np.float32,
         "BOOLEAN": np.float32
     }.get(attribute_component_type)
+
+def get_attribute_type(component_type, data_type):
+    if gltf2_io_constants.DataType.num_elements(data_type) == 1:
+        return {
+            gltf2_io_constants.ComponentType.Float: "FLOAT"
+        }[component_type]
+    elif gltf2_io_constants.DataType.num_elements(data_type) == 2:
+        return {
+            gltf2_io_constants.ComponentType.Float: "FLOAT2"
+        }[component_type]
+    elif gltf2_io_constants.DataType.num_elements(data_type) == 3:
+        return {
+            gltf2_io_constants.ComponentType.Float: "FLOAT_VECTOR"
+        }[component_type]
+    elif gltf2_io_constants.DataType.num_elements(data_type) == 4:
+        return {
+            gltf2_io_constants.ComponentType.Float: "FLOAT_COLOR",
+            gltf2_io_constants.ComponentType.UnsignedShort: "BYTE_COLOR"
+        }[component_type]
+    else:
+        pass
\ No newline at end of file
diff --git a/io_scene_gltf2/blender/imp/gltf2_blender_mesh.py b/io_scene_gltf2/blender/imp/gltf2_blender_mesh.py
index f8d21f23..bb8d51f4 100755
--- a/io_scene_gltf2/blender/imp/gltf2_blender_mesh.py
+++ b/io_scene_gltf2/blender/imp/gltf2_blender_mesh.py
@@ -9,6 +9,8 @@ from ...io.imp.gltf2_io_binary import BinaryData
 from ..com.gltf2_blender_extras import set_extras
 from .gltf2_blender_material import BlenderMaterial
 from ...io.com.gltf2_io_debug import print_console
+from ...io.com.gltf2_io_constants import DataType, ComponentType
+from ...blender.com.gltf2_blender_conversion import get_attribute_type
 from .gltf2_io_draco_compression_extension import decode_primitive
 from io_scene_gltf2.io.imp.gltf2_io_user_extensions import import_user_extensions
 from ..com.gltf2_blender_ui import gltf2_KHR_materials_variants_primitive, gltf2_KHR_materials_variants_variant, gltf2_KHR_materials_variants_default_material
@@ -75,6 +77,7 @@ def do_primitives(gltf, mesh_idx, skin_idx, mesh, ob):
     num_uvs = 0
     num_cols = 0
     num_joint_sets = 0
+    attributes = set({})
     for prim in pymesh.primitives:
         if 'POSITION' not in prim.attributes:
             continue
@@ -98,6 +101,8 @@ def do_primitives(gltf, mesh_idx, skin_idx, mesh, ob):
         while i < COLOR_MAX and ('COLOR_%d' % i) in prim.attributes: i += 1
         num_cols = max(i, num_cols)
 
+        attributes.update(set([k for k in prim.attributes if k.startswith('_')]))
+
     num_shapekeys = sum(sk_name is not None for sk_name in pymesh.shapekey_names)
 
     # -------------
@@ -129,6 +134,13 @@ def do_primitives(gltf, mesh_idx, skin_idx, mesh, ob):
         np.empty(dtype=np.float32, shape=(0,3))  # coordinate for each vert for each shapekey
         for _ in range(num_shapekeys)
     ]
+    attribute_data = []
+    for attr in attributes:
+        attribute_data.append(
+            np.empty(
+                dtype=ComponentType.to_numpy_dtype(gltf.data.accessors[prim.attributes[attr]].component_type), 
+                shape=(0, DataType.num_elements(gltf.data.accessors[prim.attributes[attr]].type)))
+                )
 
     for prim in pymesh.primitives:
         prim.num_faces = 0
@@ -230,6 +242,16 @@ def do_primitives(gltf, mesh_idx, skin_idx, mesh, ob):
                     cols = np.ones((len(indices), 4), dtype=np.float32)
                 loop_cols[col_i] = np.concatenate((loop_cols[col_i], cols))
 
+        for idx, attr in enumerate(attributes):
+            if attr in prim.attributes:
+                attr_data = BinaryData.decode_accessor(gltf, prim.attributes[attr], cache=True)
+            else:
+                attr_data = np.zeros(
+                    (len(indices), DataType.num_elements(gltf.data.accessors[prim.attributes[attr]].type)),
+                     dtype=ComponentType.to_numpy_dtype(gltf.data.accessors[prim.attributes[attr]].component_type)
+                )
+            attribute_data[idx] = np.concatenate((attribute_data[idx], attr_data))
+
     # Accessors are cached in case they are shared between primitives; clear
     # the cache now that all prims are done.
     gltf.decode_accessor_cache = {}
@@ -422,6 +444,23 @@ def do_primitives(gltf, mesh_idx, skin_idx, mesh, ob):
 
         mesh.polygons.foreach_set('material_index', material_indices)
 
+    # Custom Attributes
+    for idx, attr in enumerate(attributes):
+
+        blender_attribute_data_type = get_attribute_type(
+            gltf.data.accessors[prim.attributes[attr]].component_type,
+            gltf.data.accessors[prim.attributes[attr]].type
+        )
+
+        blender_attribute = mesh.attributes.new(attr, blender_attribute_data_type, 'POINT')
+        if DataType.num_elements(gltf.data.accessors[prim.attributes[attr]].type) == 1:
+            blender_attribute.data.foreach_set('value', attribute_data[idx].flatten())
+        elif DataType.num_elements(gltf.data.accessors[prim.attributes[attr]].type) > 1:
+            if blender_attribute_data_type in ["BYTE_COLOR", "FLOAT_COLOR"]:
+                blender_attribute.data.foreach_set('color', attribute_data[idx].flatten())
+            else:
+                blender_attribute.data.foreach_set('vector', attribute_data[idx].flatten())
+
     # ----
     # Normals