glTF importer: speedup import of vertex color

4 files changed, 62 insertions, 14 deletions

diff --git a/io_scene_gltf2/__init__.py b/io_scene_gltf2/__init__.py
index f9ad9b8d..49f6f470 100755
--- a/io_scene_gltf2/__init__.py
+++ b/io_scene_gltf2/__init__.py
@@ -15,7 +15,7 @@
 bl_info = {
     'name': 'glTF 2.0 format',
     'author': 'Julien Duroure, Norbert Nopper, Urs Hanselmann, Moritz Becher, Benjamin Schmithüsen, Jim Eckerlein, and many external contributors',
-    "version": (1, 2, 37),
+    "version": (1, 2, 38),
     'blender': (2, 82, 7),
     'location': 'File > Import-Export',
     'description': 'Import-Export as glTF 2.0',
diff --git a/io_scene_gltf2/blender/imp/gltf2_blender_gltf.py b/io_scene_gltf2/blender/imp/gltf2_blender_gltf.py
index 8e6c1950..e00e2449 100755
--- a/io_scene_gltf2/blender/imp/gltf2_blender_gltf.py
+++ b/io_scene_gltf2/blender/imp/gltf2_blender_gltf.py
@@ -24,7 +24,26 @@ class BlenderGlTF():
 
     @staticmethod
     def create(gltf):
-        """Create glTF main method."""
+        """Create glTF main method, with optional profiling"""
+        profile = bpy.app.debug_value == 102
+        if profile:
+            import cProfile, pstats, io
+            from pstats import SortKey
+            pr = cProfile.Profile()
+            pr.enable()
+            BlenderGlTF._create(gltf)
+            pr.disable()
+            s = io.StringIO()
+            sortby = SortKey.TIME
+            ps = pstats.Stats(pr, stream=s).sort_stats(sortby)
+            ps.print_stats()
+            print(s.getvalue())
+        else:
+            BlenderGlTF._create(gltf)
+
+    @staticmethod
+    def _create(gltf):
+        """Create glTF main worker method."""
         BlenderGlTF.set_convert_functions(gltf)
         BlenderGlTF.pre_compute(gltf)
         BlenderScene.create(gltf)
diff --git a/io_scene_gltf2/blender/imp/gltf2_blender_primitive.py b/io_scene_gltf2/blender/imp/gltf2_blender_primitive.py
index 07030a62..a4df6f3f 100755
--- a/io_scene_gltf2/blender/imp/gltf2_blender_primitive.py
+++ b/io_scene_gltf2/blender/imp/gltf2_blender_primitive.py
@@ -14,6 +14,8 @@
 
 import bpy
 from mathutils import Vector, Matrix
+import numpy as np
+# import time
 
 from .gltf2_blender_material import BlenderMaterial
 from ...io.imp.gltf2_io_binary import BinaryData
@@ -194,20 +196,27 @@ class BlenderPrimitive():
             layer_name = 'Col' if set_num == 0 else 'Col.%03d' % set_num
             layer = BlenderPrimitive.get_layer(bme.loops.layers.color, layer_name)
 
-            colors = BinaryData.get_data_from_accessor(gltf, attributes['COLOR_%d' % set_num], cache=True)
+            # colors is a 2d array: [N][3 or 4]
+            gltf_attr_name = 'COLOR_%d' % set_num
+            colors_raw = BinaryData.get_data_from_accessor(gltf, attributes[gltf_attr_name], cache=True)
+            colors = np.array(colors_raw, dtype=np.float32)
 
             is_rgba = len(colors[0]) == 4
-
+            if not is_rgba:
+                # RGB -> RGBA
+                ones = np.ones((colors.shape[0], 1))
+                colors = np.concatenate((colors, ones), axis=1) # add alpha channel
+
+            srgb_colors = color_linear_to_srgb(colors)
+            # t = time.perf_counter()
+            # This needs to be a tight loop because it runs over all vertices,
+            # which is why this code looks a little odd.
             for bidx, pidx in vert_idxs:
-                color = colors[pidx]
-                col = (
-                    color_linear_to_srgb(color[0]),
-                    color_linear_to_srgb(color[1]),
-                    color_linear_to_srgb(color[2]),
-                    color[3] if is_rgba else 1.0,
-                )
+                color = srgb_colors[pidx]
+                col = (color[0], color[1], color[2], color[3]) # fastest this way
                 for loop in bme_verts[bidx].link_loops:
                     loop[layer] = col
+            # print(f'store colors: {time.perf_counter() - t}')
 
             set_num += 1
 
diff --git a/io_scene_gltf2/io/com/gltf2_io_color_management.py b/io_scene_gltf2/io/com/gltf2_io_color_management.py
index 56b3a246..58dc3a27 100644
--- a/io_scene_gltf2/io/com/gltf2_io_color_management.py
+++ b/io_scene_gltf2/io/com/gltf2_io_color_management.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
 
 def color_srgb_to_scene_linear(c):
     """
@@ -29,8 +30,27 @@ def color_linear_to_srgb(c):
     Convert from linear to sRGB color space.
 
     Source: Cycles addon implementation, node_color.h.
+    c may be a single color value or an array.
+
+    If c's last dimension is 4, it's assumed to be RGBA and the
+    alpha channel is not converted.
     """
-    if c < 0.0031308:
-        return 0.0 if c < 0.0 else c * 12.92
+    if type(c) in (list, np.ndarray):
+        colors = np.array(c, np.float32) if type(c) == list else c
+        if  colors.ndim > 1 and colors.shape[-1] == 4:
+            colors_noa = colors[..., 0:3] # only process RGB for speed
+        else:
+            colors_noa = colors
+        not_small = colors_noa >= 0.0031308
+        small_result = np.where(colors_noa < 0.0, 0.0, colors_noa * 12.92)
+        large_result = 1.055 * np.power(colors_noa, 1.0 / 2.4, where=not_small) - 0.055
+        result = np.where(not_small, large_result, small_result)
+        if  colors.ndim > 1 and colors.shape[-1] == 4:
+            # copy alpha from original
+            result = np.concatenate((result, colors[..., 3, np.newaxis]), axis=-1)
+        return result
     else:
-        return 1.055 * pow(c, 1.0 / 2.4) - 0.055
+        if c < 0.0031308:
+            return 0.0 if c < 0.0 else c * 12.92
+        else:
+            return 1.055 * pow(c, 1.0 / 2.4) - 0.055