pose_library/pose_creation.py

   1 # SPDX-FileCopyrightText: 2021-2023 Blender Foundation
   2 #
   3 # SPDX-License-Identifier: GPL-2.0-or-later
   4
   5 """
   6 Pose Library - creation functions.
   7 """
   8
   9 import dataclasses
  10 import functools
  11 import re
  12
  13 from typing import Optional, FrozenSet, Set, Union, Iterable, cast
  14
  15 if "functions" not in locals():
  16     from . import asset_browser, functions
  17 else:
  18     import importlib
  19
  20     asset_browser = importlib.reload(asset_browser)
  21     functions = importlib.reload(functions)
  22
  23 import bpy
  24 from bpy.types import (
  25     Action,
  26     Bone,
  27     Context,
  28     FCurve,
  29     Keyframe,
  30 )
  31
  32 FCurveValue = Union[float, int]
  33
  34 pose_bone_re = re.compile(r'pose.bones\["([^"]+)"\]')
  35 """RegExp for matching FCurve data paths."""
  36
  37
  38 @dataclasses.dataclass(unsafe_hash=True, frozen=True)
  39 class PoseCreationParams:
  40     armature_ob: bpy.types.Object
  41     src_action: Optional[Action]
  42     src_frame_nr: float
  43     bone_names: FrozenSet[str]
  44     new_asset_name: str
  45
  46
  47 class UnresolvablePathError(ValueError):
  48     """Raised when a data_path cannot be resolved to a current value."""
  49
  50
  51 @dataclasses.dataclass(unsafe_hash=True)
  52 class PoseActionCreator:
  53     """Create an Action that's suitable for marking as Asset.
  54
  55     Does not mark as asset yet, nor does it add asset metadata.
  56     """
  57
  58     params: PoseCreationParams
  59
  60     # These were taken from Blender's Action baking code in `anim_utils.py`.
  61     # Items are (name, array_length) tuples.
  62     _bbone_props = [
  63         ("bbone_curveinx", None),
  64         ("bbone_curveoutx", None),
  65         ("bbone_curveinz", None),
  66         ("bbone_curveoutz", None),
  67         ("bbone_rollin", None),
  68         ("bbone_rollout", None),
  69         ("bbone_scalein", 3),
  70         ("bbone_scaleout", 3),
  71         ("bbone_easein", None),
  72         ("bbone_easeout", None),
  73     ]
  74
  75     def create(self) -> Optional[Action]:
  76         """Create a single-frame Action containing only the given bones, or None if no anim data was found."""
  77
  78         try:
  79             dst_action = self._create_new_action()
  80             self._store_pose(dst_action)
  81         finally:
  82             # Prevent next instantiations of this class from reusing pointers to
  83             # bones. They may not be valid by then any more.
  84             self._find_bone.cache_clear()
  85
  86         if len(dst_action.fcurves) == 0:
  87             bpy.data.actions.remove(dst_action)
  88             return None
  89
  90         return dst_action
  91
  92     def _create_new_action(self) -> Action:
  93         dst_action = bpy.data.actions.new(self.params.new_asset_name)
  94         if self.params.src_action:
  95             dst_action.id_root = self.params.src_action.id_root
  96         dst_action.user_clear()  # actions.new() sets users=1, but marking as asset also increments user count.
  97         return dst_action
  98
  99     def _store_pose(self, dst_action: Action) -> None:
 100         """Store the current pose into the given action."""
 101         self._store_bone_pose_parameters(dst_action)
 102         self._store_animated_parameters(dst_action)
 103
 104     def _store_bone_pose_parameters(self, dst_action: Action) -> None:
 105         """Store loc/rot/scale/bbone values in the Action."""
 106
 107         for bone_name in sorted(self.params.bone_names):
 108             self._store_location(dst_action, bone_name)
 109             self._store_rotation(dst_action, bone_name)
 110             self._store_scale(dst_action, bone_name)
 111             self._store_bbone(dst_action, bone_name)
 112
 113     def _store_animated_parameters(self, dst_action: Action) -> None:
 114         """Store the current value of any animated bone properties."""
 115         if self.params.src_action is None:
 116             return
 117
 118         armature_ob = self.params.armature_ob
 119         for fcurve in self.params.src_action.fcurves:
 120             match = pose_bone_re.match(fcurve.data_path)
 121             if not match:
 122                 # Not animating a bone property.
 123                 continue
 124
 125             bone_name = match.group(1)
 126             if bone_name not in self.params.bone_names:
 127                 # Bone is not our export set.
 128                 continue
 129
 130             if dst_action.fcurves.find(fcurve.data_path, index=fcurve.array_index):
 131                 # This property is already handled by a previous _store_xxx() call.
 132                 continue
 133
 134             # Only include in the pose if there is a key on this frame.
 135             if not self._has_key_on_frame(fcurve):
 136                 continue
 137
 138             try:
 139                 value = self._current_value(armature_ob, fcurve.data_path, fcurve.array_index)
 140             except UnresolvablePathError:
 141                 # A once-animated property no longer exists.
 142                 continue
 143
 144             dst_fcurve = dst_action.fcurves.new(fcurve.data_path, index=fcurve.array_index, action_group=bone_name)
 145             dst_fcurve.keyframe_points.insert(self.params.src_frame_nr, value=value)
 146             dst_fcurve.update()
 147
 148     def _store_location(self, dst_action: Action, bone_name: str) -> None:
 149         """Store bone location."""
 150         self._store_bone_array(dst_action, bone_name, "location", 3)
 151
 152     def _store_rotation(self, dst_action: Action, bone_name: str) -> None:
 153         """Store bone rotation given current rotation mode."""
 154         bone = self._find_bone(bone_name)
 155         if bone.rotation_mode == "QUATERNION":
 156             self._store_bone_array(dst_action, bone_name, "rotation_quaternion", 4)
 157         elif bone.rotation_mode == "AXIS_ANGLE":
 158             self._store_bone_array(dst_action, bone_name, "rotation_axis_angle", 4)
 159         else:
 160             self._store_bone_array(dst_action, bone_name, "rotation_euler", 3)
 161
 162     def _store_scale(self, dst_action: Action, bone_name: str) -> None:
 163         """Store bone scale."""
 164         self._store_bone_array(dst_action, bone_name, "scale", 3)
 165
 166     def _store_bbone(self, dst_action: Action, bone_name: str) -> None:
 167         """Store bendy-bone parameters."""
 168         for prop_name, array_length in self._bbone_props:
 169             if array_length:
 170                 self._store_bone_array(dst_action, bone_name, prop_name, array_length)
 171             else:
 172                 self._store_bone_property(dst_action, bone_name, prop_name)
 173
 174     def _store_bone_array(self, dst_action: Action, bone_name: str, property_name: str, array_length: int) -> None:
 175         """Store all elements of an array property."""
 176         for array_index in range(array_length):
 177             self._store_bone_property(dst_action, bone_name, property_name, array_index)
 178
 179     def _store_bone_property(
 180         self,
 181         dst_action: Action,
 182         bone_name: str,
 183         property_path: str,
 184         array_index: int = -1,
 185     ) -> None:
 186         """Store the current value of a single bone property."""
 187
 188         bone = self._find_bone(bone_name)
 189         value = self._current_value(bone, property_path, array_index)
 190
 191         # Get the full 'pose.bones["bone_name"].blablabla' path suitable for FCurves.
 192         rna_path = bone.path_from_id(property_path)
 193
 194         fcurve: Optional[FCurve] = dst_action.fcurves.find(rna_path, index=array_index)
 195         if fcurve is None:
 196             fcurve = dst_action.fcurves.new(rna_path, index=array_index, action_group=bone_name)
 197
 198         fcurve.keyframe_points.insert(self.params.src_frame_nr, value=value)
 199         fcurve.update()
 200
 201     @classmethod
 202     def _current_value(cls, datablock: bpy.types.ID, data_path: str, array_index: int) -> FCurveValue:
 203         """Resolve an RNA path + array index to an actual value."""
 204         value_or_array = cls._path_resolve(datablock, data_path)
 205
 206         # Both indices -1 and 0 are used for non-array properties.
 207         # -1 cannot be used in arrays, whereas 0 can be used in both arrays and non-arrays.
 208
 209         if array_index == -1:
 210             return cast(FCurveValue, value_or_array)
 211
 212         if array_index == 0:
 213             value_or_array = cls._path_resolve(datablock, data_path)
 214             try:
 215                 # MyPy doesn't understand this try/except is to determine the type.
 216                 value = value_or_array[array_index]  # type: ignore
 217             except TypeError:
 218                 # Not an array after all.
 219                 return cast(FCurveValue, value_or_array)
 220             return cast(FCurveValue, value)
 221
 222         # MyPy doesn't understand that array_index>0 implies this is indexable.
 223         return cast(FCurveValue, value_or_array[array_index])  # type: ignore
 224
 225     @staticmethod
 226     def _path_resolve(datablock: bpy.types.ID, data_path: str) -> Union[FCurveValue, Iterable[FCurveValue]]:
 227         """Wrapper for datablock.path_resolve(data_path).
 228
 229         Raise UnresolvablePathError when the path cannot be resolved.
 230         This is easier to deal with upstream than the generic ValueError raised
 231         by Blender.
 232         """
 233         try:
 234             return datablock.path_resolve(data_path)  # type: ignore
 235         except ValueError as ex:
 236             raise UnresolvablePathError(str(ex)) from ex
 237
 238     @functools.lru_cache(maxsize=1024)
 239     def _find_bone(self, bone_name: str) -> Bone:
 240         """Find a bone by name.
 241
 242         Assumes the named bone exists, as the bones this class handles comes
 243         from the user's selection, and you can't select a non-existent bone.
 244         """
 245
 246         bone: Bone = self.params.armature_ob.pose.bones[bone_name]
 247         return bone
 248
 249     def _has_key_on_frame(self, fcurve: FCurve) -> bool:
 250         """Return True iff the FCurve has a key on the source frame."""
 251
 252         points = fcurve.keyframe_points
 253         if not points:
 254             return False
 255
 256         frame_to_find = self.params.src_frame_nr
 257         margin = 0.001
 258         high = len(points) - 1
 259         low = 0
 260         while low <= high:
 261             mid = (high + low) // 2
 262             diff = points[mid].co.x - frame_to_find
 263             if abs(diff) < margin:
 264                 return True
 265             if diff < 0:
 266                 # Frame to find is bigger than the current middle.
 267                 low = mid + 1
 268             else:
 269                 # Frame to find is smaller than the current middle
 270                 high = mid - 1
 271         return False
 272
 273
 274 def create_pose_asset(
 275     params: PoseCreationParams,
 276 ) -> Optional[Action]:
 277     """Create a single-frame Action containing only the pose of the given bones.
 278
 279     DOES mark as asset, DOES NOT configure asset metadata.
 280     """
 281
 282     creator = PoseActionCreator(params)
 283     pose_action = creator.create()
 284     if pose_action is None:
 285         return None
 286
 287     pose_action.asset_mark()
 288     pose_action.asset_generate_preview()
 289     return pose_action
 290
 291
 292 def create_pose_asset_from_context(context: Context, new_asset_name: str) -> Optional[Action]:
 293     """Create Action asset from active object & selected bones."""
 294
 295     bones = context.selected_pose_bones_from_active_object
 296     bone_names = {bone.name for bone in bones}
 297
 298     params = PoseCreationParams(
 299         context.object,
 300         getattr(context.object.animation_data, "action", None),
 301         context.scene.frame_current,
 302         frozenset(bone_names),
 303         new_asset_name,
 304     )
 305
 306     return create_pose_asset(params)
 307
 308
 309 def copy_fcurves(
 310     dst_action: Action,
 311     src_action: Action,
 312     src_frame_nr: float,
 313     bone_names: Set[str],
 314 ) -> int:
 315     """Copy FCurves, returning number of curves copied."""
 316     num_fcurves_copied = 0
 317     for fcurve in src_action.fcurves:
 318         match = pose_bone_re.match(fcurve.data_path)
 319         if not match:
 320             continue
 321
 322         bone_name = match.group(1)
 323         if bone_name not in bone_names:
 324             continue
 325
 326         # Check if there is a keyframe on this frame.
 327         keyframe = find_keyframe(fcurve, src_frame_nr)
 328         if keyframe is None:
 329             continue
 330         create_single_key_fcurve(dst_action, fcurve, keyframe)
 331         num_fcurves_copied += 1
 332     return num_fcurves_copied
 333
 334
 335 def create_single_key_fcurve(dst_action: Action, src_fcurve: FCurve, src_keyframe: Keyframe) -> FCurve:
 336     """Create a copy of the source FCurve, but only for the given keyframe.
 337
 338     Returns a new FCurve with just one keyframe.
 339     """
 340
 341     dst_fcurve = copy_fcurve_without_keys(dst_action, src_fcurve)
 342     copy_keyframe(dst_fcurve, src_keyframe)
 343     return dst_fcurve
 344
 345
 346 def copy_fcurve_without_keys(dst_action: Action, src_fcurve: FCurve) -> FCurve:
 347     """Create a new FCurve and copy some properties."""
 348
 349     src_group_name = src_fcurve.group.name if src_fcurve.group else ""
 350     dst_fcurve = dst_action.fcurves.new(src_fcurve.data_path, index=src_fcurve.array_index, action_group=src_group_name)
 351     for propname in {"auto_smoothing", "color", "color_mode", "extrapolation"}:
 352         setattr(dst_fcurve, propname, getattr(src_fcurve, propname))
 353     return dst_fcurve
 354
 355
 356 def copy_keyframe(dst_fcurve: FCurve, src_keyframe: Keyframe) -> Keyframe:
 357     """Copy a keyframe from one FCurve to the other."""
 358
 359     dst_keyframe = dst_fcurve.keyframe_points.insert(
 360         src_keyframe.co.x, src_keyframe.co.y, options={'FAST'}, keyframe_type=src_keyframe.type
 361     )
 362
 363     for propname in {
 364         "amplitude",
 365         "back",
 366         "easing",
 367         "handle_left",
 368         "handle_left_type",
 369         "handle_right",
 370         "handle_right_type",
 371         "interpolation",
 372         "period",
 373     }:
 374         setattr(dst_keyframe, propname, getattr(src_keyframe, propname))
 375     dst_fcurve.update()
 376     return dst_keyframe
 377
 378
 379 def find_keyframe(fcurve: FCurve, frame: float) -> Optional[Keyframe]:
 380     # Binary search adapted from https://pythonguides.com/python-binary-search/
 381     keyframes = fcurve.keyframe_points
 382     low = 0
 383     high = len(keyframes) - 1
 384     mid = 0
 385
 386     # Accept any keyframe that's within 'epsilon' of the requested frame.
 387     # This should account for rounding errors and the likes.
 388     epsilon = 1e-4
 389     frame_lowerbound = frame - epsilon
 390     frame_upperbound = frame + epsilon
 391     while low <= high:
 392         mid = (high + low) // 2
 393         keyframe = keyframes[mid]
 394         if keyframe.co.x < frame_lowerbound:
 395             low = mid + 1
 396         elif keyframe.co.x > frame_upperbound:
 397             high = mid - 1
 398         else:
 399             return keyframe
 400     return None
 401
 402
 403 def assign_from_asset_browser(asset: Action, asset_browser_area: bpy.types.Area) -> None:
 404     """Assign some things from the asset browser to the asset.
 405
 406     This sets the current catalog ID, and in the future could include tags
 407     from the active dynamic catalog, etc.
 408     """
 409
 410     cat_id = asset_browser.active_catalog_id(asset_browser_area)
 411     asset.asset_data.catalog_id = cat_id