From d7b591d242923a416cb40db34d7264507a258cc7 Mon Sep 17 00:00:00 2001 From: Thomas Barlow Date: Sat, 23 Dec 2023 05:17:26 +0000 Subject: [PATCH 1/3] Speed up shape key foreach_get/set access using pointers --- io_scene_fbx/export_fbx_bin.py | 4 +- io_scene_fbx/fbx_utils.py | 120 +++++++++++++++++++++++++++++++++ io_scene_fbx/import_fbx.py | 3 +- 3 files changed, 124 insertions(+), 3 deletions(-) diff --git a/io_scene_fbx/export_fbx_bin.py b/io_scene_fbx/export_fbx_bin.py index 2b1c58de0..b2941ed8f 100644 --- a/io_scene_fbx/export_fbx_bin.py +++ b/io_scene_fbx/export_fbx_bin.py @@ -48,7 +48,7 @@ from .fbx_utils import ( PerfMon, units_blender_to_fbx_factor, units_convertor, units_convertor_iter, matrix4_to_array, similar_values, shape_difference_exclude_similar, astype_view_signedness, fast_first_axis_unique, - fast_first_axis_flat, + fast_first_axis_flat, fast_mesh_shape_key_co_foreach_get, # Attribute helpers. MESH_ATTRIBUTE_CORNER_EDGE, MESH_ATTRIBUTE_SHARP_EDGE, MESH_ATTRIBUTE_EDGE_VERTS, MESH_ATTRIBUTE_CORNER_VERT, MESH_ATTRIBUTE_SHARP_FACE, MESH_ATTRIBUTE_POSITION, MESH_ATTRIBUTE_MATERIAL_INDEX, @@ -2753,7 +2753,7 @@ def fbx_data_from_scene(scene, depsgraph, settings): _cos = MESH_ATTRIBUTE_POSITION.to_ndarray(me.attributes) else: _cos = np.empty(len(me.vertices) * 3, dtype=co_bl_dtype) - shape_key.data.foreach_get("co", _cos) + fast_mesh_shape_key_co_foreach_get(shape_key, _cos) return vcos_transformed(_cos, geom_mat_co, co_fbx_dtype) for shape in me.shape_keys.key_blocks[1:]: diff --git a/io_scene_fbx/fbx_utils.py b/io_scene_fbx/fbx_utils.py index 1b8a08ca3..feb050917 100644 --- a/io_scene_fbx/fbx_utils.py +++ b/io_scene_fbx/fbx_utils.py @@ -10,6 +10,7 @@ from collections import namedtuple from collections.abc import Iterable from itertools import zip_longest, chain from dataclasses import dataclass, field +from types import SimpleNamespace from typing import Callable import numpy as np @@ -659,6 +660,125 @@ def expand_shape_key_range(shape_key, value_to_fit): return True +def _shape_key_co_memory_as_ndarray(shape_key, do_check=True): + """ + ShapeKey.data elements have a dynamic typing based on the Object the shape keys are attached to, which makes them + slower to access with foreach_get. For 'MESH' type Objects, the elements are always ShapeKeyPoint type and their + data is stored contiguously, meaning an array can be constructed from the pointer to the first ShapeKeyPoint + element. + + Creating an array from a pointer is inherently unsafe, so this function does a number of checks to make it safer. + """ + if do_check and not _fast_mesh_shape_key_co_check(): + return None + shape_data = shape_key.data + num_co = len(shape_data) + + if num_co < 2: + # At least 2 elements are required to check memory size. + return None + + co_dtype = np.dtype(np.single) + + start_address = shape_data[0].as_pointer() + last_element_start_address = shape_data[-1].as_pointer() + + memory_length_minus_one_item = last_element_start_address - start_address + + expected_element_size = co_dtype.itemsize * 3 + expected_memory_length = (num_co - 1) * expected_element_size + + if memory_length_minus_one_item == expected_memory_length: + # Use NumPy's array interface protocol to construct an array from the pointer. + array_interface_holder = SimpleNamespace( + __array_interface__=dict( + shape=(num_co * 3,), + typestr=co_dtype.str, + data=(start_address, False), # False for writable + version=3, + ) + ) + return np.asarray(array_interface_holder) + else: + return None + + +# Initially set to None +_USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET = None + + +def _fast_mesh_shape_key_co_check(): + global _USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET + if _USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET is not None: + # The check has already been run and the result has been stored in _USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET. + return _USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET + + tmp_mesh = None + tmp_object = None + try: + tmp_mesh = bpy.data.meshes.new("") + num_co = 100 + tmp_mesh.vertices.add(num_co) + # An Object is needed to add/remove shape keys from a Mesh. + tmp_object = bpy.data.objects.new("", tmp_mesh) + shape_key = tmp_object.shape_key_add(name="") + shape_data = shape_key.data + + if shape_key.bl_rna.properties["data"].fixed_type == shape_data[0].bl_rna: + # The shape key "data" collection is no longer dynamically typed and foreach_get/set should be fast enough. + _USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET = False + return False + + co_dtype = np.dtype(np.single) + + # Fill the shape key with some data. + shape_data.foreach_set("co", np.arange(3 * num_co, dtype=co_dtype)) + + # The check is this function, so explicitly don't do the check. + co_memory_as_array = _shape_key_co_memory_as_ndarray(shape_key, do_check=False) + if co_memory_as_array is not None: + # Immediately make a copy in case the `foreach_get` afterwards can cause the memory to be reallocated. + co_array_from_memory = co_memory_as_array.copy() + del co_memory_as_array + # Check that the array created from the pointer has the exact same contents + # as using foreach_get. + co_array_check = np.empty(num_co * 3, dtype=co_dtype) + shape_data.foreach_get("co", co_array_check) + if np.array_equal(co_array_check, co_array_from_memory, equal_nan=True): + _USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET = True + return True + + # Something didn't work. + _USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET = False + return False + finally: + # Clean up temporary objects. + if tmp_object is not None: + tmp_object.shape_key_clear() + bpy.data.objects.remove(tmp_object) + if tmp_mesh is not None: + bpy.data.meshes.remove(tmp_mesh) + + +def fast_mesh_shape_key_co_foreach_get(shape_key, seq): + co_memory_as_array = _shape_key_co_memory_as_ndarray(shape_key) + if co_memory_as_array is not None: + seq[:] = co_memory_as_array + else: + shape_key.data.foreach_get("co", seq) + + +def fast_mesh_shape_key_co_foreach_set(shape_key, seq): + co_memory_as_array = _shape_key_co_memory_as_ndarray(shape_key) + if co_memory_as_array is not None: + co_memory_as_array[:] = seq + # Memory has been set directly, so call .update() manually. + # TODO: Not sure if this is required + shape_key.data.update() + else: + shape_key.data.foreach_set("co", seq) + + # ##### Attribute utils. ##### AttributeDataTypeInfo = namedtuple("AttributeDataTypeInfo", ["dtype", "foreach_attribute", "item_size"]) _attribute_data_type_info_lookup = { diff --git a/io_scene_fbx/import_fbx.py b/io_scene_fbx/import_fbx.py index f306e9039..a6c6e7830 100644 --- a/io_scene_fbx/import_fbx.py +++ b/io_scene_fbx/import_fbx.py @@ -51,6 +51,7 @@ from .fbx_utils import ( FBX_KTIME_V7, FBX_KTIME_V8, FBX_TIMECODE_DEFINITION_TO_KTIME_PER_SECOND, + fast_mesh_shape_key_co_foreach_set, ) LINEAR_INTERPOLATION_VALUE = bpy.types.Keyframe.bl_rna.properties['interpolation'].enum_items['LINEAR'].value @@ -2002,7 +2003,7 @@ def blen_read_shapes(fbx_tmpl, fbx_data, objects, me, scene): if dvcos.any(): shape_cos = me_vcos_vector_view.copy() shape_cos[indices] += dvcos - kb.data.foreach_set("co", shape_cos.ravel()) + fast_mesh_shape_key_co_foreach_set(kb, shape_cos.ravel()) shape_key_values_in_range &= expand_shape_key_range(kb, weight) -- 2.30.2 From 562c371eacf961bbce7c8ee82cc27d825a04561f Mon Sep 17 00:00:00 2001 From: Thomas Barlow Date: Fri, 12 Jan 2024 22:23:19 +0000 Subject: [PATCH 2/3] Index the shape key data memory directly This skips the copying steps. --- io_scene_fbx/fbx_utils.py | 22 +++++++++++++++------- io_scene_fbx/import_fbx.py | 6 ++---- 2 files changed, 17 insertions(+), 11 deletions(-) diff --git a/io_scene_fbx/fbx_utils.py b/io_scene_fbx/fbx_utils.py index feb050917..6c49e8a8a 100644 --- a/io_scene_fbx/fbx_utils.py +++ b/io_scene_fbx/fbx_utils.py @@ -768,15 +768,23 @@ def fast_mesh_shape_key_co_foreach_get(shape_key, seq): shape_key.data.foreach_get("co", seq) -def fast_mesh_shape_key_co_foreach_set(shape_key, seq): - co_memory_as_array = _shape_key_co_memory_as_ndarray(shape_key) +def fast_mesh_shape_key_dvcos_foreach_set(new_shape_key, dvcos, indices, mesh_positions_fallback_vector_view): + """ + FBX Shape key data are sparse vectors relative to the mesh, unlike Blender shape keys which are coordinates. + + The newly created shape keys must have been created with `from_mix=False`, so that they match the mesh positions. + """ + co_memory_as_array = _shape_key_co_memory_as_ndarray(new_shape_key) if co_memory_as_array is not None: - co_memory_as_array[:] = seq - # Memory has been set directly, so call .update() manually. - # TODO: Not sure if this is required - shape_key.data.update() + co_memory_as_array_vector_view = co_memory_as_array.view() + co_memory_as_array_vector_view.shape = (-1, 3) + co_memory_as_array_vector_view[indices] += dvcos + # Memory has been set directly, so call .update(). + new_shape_key.data.update() else: - shape_key.data.foreach_set("co", seq) + shape_cos = mesh_positions_fallback_vector_view.copy() + shape_cos[indices] += dvcos + new_shape_key.data.foreach_set("co", shape_cos.ravel()) # ##### Attribute utils. ##### diff --git a/io_scene_fbx/import_fbx.py b/io_scene_fbx/import_fbx.py index a6c6e7830..5e52e6237 100644 --- a/io_scene_fbx/import_fbx.py +++ b/io_scene_fbx/import_fbx.py @@ -51,7 +51,7 @@ from .fbx_utils import ( FBX_KTIME_V7, FBX_KTIME_V8, FBX_TIMECODE_DEFINITION_TO_KTIME_PER_SECOND, - fast_mesh_shape_key_co_foreach_set, + fast_mesh_shape_key_dvcos_foreach_set, ) LINEAR_INTERPOLATION_VALUE = bpy.types.Keyframe.bl_rna.properties['interpolation'].enum_items['LINEAR'].value @@ -2001,9 +2001,7 @@ def blen_read_shapes(fbx_tmpl, fbx_data, objects, me, scene): # Only need to set the shape key co if there are any non-zero dvcos. if dvcos.any(): - shape_cos = me_vcos_vector_view.copy() - shape_cos[indices] += dvcos - fast_mesh_shape_key_co_foreach_set(kb, shape_cos.ravel()) + fast_mesh_shape_key_dvcos_foreach_set(kb, dvcos, indices, me_vcos_vector_view) shape_key_values_in_range &= expand_shape_key_range(kb, weight) -- 2.30.2 From 18c32f693fb94fce3a7bae3a00ea57ac32876ee0 Mon Sep 17 00:00:00 2001 From: Thomas Barlow Date: Sat, 13 Jan 2024 22:40:39 +0000 Subject: [PATCH 3/3] More comments --- io_scene_fbx/fbx_utils.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/io_scene_fbx/fbx_utils.py b/io_scene_fbx/fbx_utils.py index 6c49e8a8a..ad90fb089 100644 --- a/io_scene_fbx/fbx_utils.py +++ b/io_scene_fbx/fbx_utils.py @@ -713,6 +713,8 @@ def _fast_mesh_shape_key_co_check(): # The check has already been run and the result has been stored in _USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET. return _USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET + # Check that accessing a shape key's data through its pointer works, by creating a temporary mesh and adding a shape + # key to it. tmp_mesh = None tmp_object = None try: @@ -737,14 +739,15 @@ def _fast_mesh_shape_key_co_check(): # The check is this function, so explicitly don't do the check. co_memory_as_array = _shape_key_co_memory_as_ndarray(shape_key, do_check=False) if co_memory_as_array is not None: - # Immediately make a copy in case the `foreach_get` afterwards can cause the memory to be reallocated. + # Immediately make a copy in the unlikely case the `foreach_get` call afterward can cause the memory to be + # reallocated. co_array_from_memory = co_memory_as_array.copy() del co_memory_as_array # Check that the array created from the pointer has the exact same contents # as using foreach_get. - co_array_check = np.empty(num_co * 3, dtype=co_dtype) - shape_data.foreach_get("co", co_array_check) - if np.array_equal(co_array_check, co_array_from_memory, equal_nan=True): + co_array_from_foreach_get = np.empty(num_co * 3, dtype=co_dtype) + shape_data.foreach_get("co", co_array_from_foreach_get) + if np.array_equal(co_array_from_foreach_get, co_array_from_memory, equal_nan=True): _USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET = True return True @@ -761,6 +764,9 @@ def _fast_mesh_shape_key_co_check(): def fast_mesh_shape_key_co_foreach_get(shape_key, seq): + """ + Replacement for ShapeKey.data.foreach_get that accesses the shape key data's memory directly if possible. + """ co_memory_as_array = _shape_key_co_memory_as_ndarray(shape_key) if co_memory_as_array is not None: seq[:] = co_memory_as_array @@ -770,9 +776,10 @@ def fast_mesh_shape_key_co_foreach_get(shape_key, seq): def fast_mesh_shape_key_dvcos_foreach_set(new_shape_key, dvcos, indices, mesh_positions_fallback_vector_view): """ - FBX Shape key data are sparse vectors relative to the mesh, unlike Blender shape keys which are coordinates. + Apply sparse FBX shape key vectors to a newly created shape key. - The newly created shape keys must have been created with `from_mix=False`, so that they match the mesh positions. + The newly created shape key must have been created with `from_mix=False`, so that its coordinates match the mesh + positions. """ co_memory_as_array = _shape_key_co_memory_as_ndarray(new_shape_key) if co_memory_as_array is not None: @@ -782,6 +789,8 @@ def fast_mesh_shape_key_dvcos_foreach_set(new_shape_key, dvcos, indices, mesh_po # Memory has been set directly, so call .update(). new_shape_key.data.update() else: + # As a fallback, copy the mesh positions, apply the sparse vectors to that copy and then set all the shape key + # coordinates to the modified copy of the mesh positions. shape_cos = mesh_positions_fallback_vector_view.copy() shape_cos[indices] += dvcos new_shape_key.data.foreach_set("co", shape_cos.ravel()) -- 2.30.2