pose_library/pose_creation.py

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