animation import export class in blender 2.79

ArmatureAnimationIO.py

#ArmatureAnimationIO.py
#tested in blender 2.79
#please modify/use at your own risk

import bpy


###for animation io
import os
import re
import json
import logging
from collections import namedtuple
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.DEBUG) #without this info logs wouldnt show in console
###

"""
import imp
testTool = imp.load_source("tool","/Users/Nathaniel/Documents/src_blender/python/animationTools/ArmatureAnimationIO.py") #change to path to python file

animIO = testTool.ArmatureAnimationIO()

#writing animation
animIO.toJson("Armature","/Users/Nathaniel/Documents/src_blender/python/snippets/tmp/anim_test.json")

#importing animation
#animIO.fromJson("Armature","/Users/Nathaniel/Documents/src_blender/python/snippets/tmp/anim_test.json")

"""


class ArmatureAnimationIO(object):
    """handles armature animation import and export
    
    example usage:
    animIO = ArmatureAnimationIO()
    
    #writing animation
    animIO.toJson("Armature","/Users/Nathaniel/Documents/src_blender/python/snippets/tmp/anim_test.json")
    
    #importing animation
    #animIO.fromJson("Armature","/Users/Nathaniel/Documents/src_blender/python/snippets/tmp/anim_test.json")
    """
    def __init__(self):
        self._animatedKey = "animated" #what key of dictionary holds animation data
        self._staticKey = "static"
        
    def toJson(self, armature, filePath):
        """export animation to json
        #Armature{
        #   animated{
        #       5:{ "bone1":channel:array_index:{bone data},
        #           "bone10":channel:array_index:{bone data}
        #            }
        #       10:{ "bone3":channel:array_index:{bone data}
        #           }
        #       
        #   },
        #   static{
        #       1:{"bone50":{bone data}, #all static can be on first frame
        #           "bone30":{bone data}
        #         }
        #   }
        #}
                
        todo: support export of a specific action
        
        @param armature (str) armature name
        @param filePath (str) output file path - directory must exist
        """
        
        if not os.path.isdir(os.path.dirname(filePath)):
            logger.warning("requires directory to exist for {}".format(filePath))
            return False
        
        if not armature in bpy.data.objects:
            logger.warning("requires armature name to exist for {}".format(armature))
            return False
            
        result = {}
        animatedKey = self._animatedKey#"animated"
        staticKey = self._staticKey#"static"
        result[armature] = {animatedKey:{}, staticKey:{}}
        animatedBones = []
        
        allFcurves = bpy.data.objects[armature].animation_data.action.fcurves
        decimals = 5 #round number decimals
        for fcurve in allFcurves:
            boneDataList = []
            keyframePoints = fcurve.keyframe_points
            for keyframePoint in keyframePoints:
                frame = keyframePoint.co[0]
                value = keyframePoint.co[1]
                value = round(value, decimals)
                interpolation = keyframePoint.interpolation
                handle_left = tuple(keyframePoint.handle_left) #need to convert Vector to tuple
                handle_left = [round(x, decimals) for x in handle_left]
                handle_left_type = keyframePoint.handle_left_type
                handle_right = tuple(keyframePoint.handle_right)
                handle_right = [round(x, decimals) for x in handle_right]
                handle_right_type = keyframePoint.handle_right_type
                array_index = fcurve.array_index
                
                ###it should handle bendybone attributes if they exist. the channel tells which channel to animate
                
                ###
                
                logger.info("frame {frame} value {value} interpolation {interpolation} handle_left {handle_left} handle_right {handle_right} handle_left_type {handle_left_type} handle_right_type {handle_right_type}".format(frame=frame, 
                                                                value=value,
                                                                interpolation=interpolation,
                                                                handle_left=handle_left,
                                                                handle_left_type=handle_left_type,
                                                                handle_right=handle_right,
                                                                handle_right_type=handle_right_type
                                                                ))
                #get bone name for fcurve
                matchGroups = re.match(r'pose.bones\["(.*)"\].(.*)', fcurve.data_path)
                bone = ''
                channel = ''
                if matchGroups:
                    bone = matchGroups.groups(0)[0]
                    channel = matchGroups.groups(0)[1]
                
                animatedBones.append(bone) #keep track of animated bones for handling static bones later
                
                boneData = dict(value=value,
                                channel=channel,
                                array_index=array_index, #ex: 0 with channel location meats location in X axis 
                                interpolation=interpolation,
                                handle_left=handle_left,
                                handle_left_type=handle_left_type,
                                handle_right=handle_right,
                                handle_right_type=handle_right_type)
                
                
                #doing the saving of dictionary
                #result[armature][animatedKey].setdefault(frame, {})[bone]=boneData #using setdefault so can create an empty dictionary if no key added yet
                result[armature][animatedKey].setdefault(frame, {}).setdefault(bone, {}).setdefault(channel, {}).setdefault(array_index,{})
                result[armature][animatedKey][frame][bone][channel][array_index] = boneData
                         
            #end keyframe points loop
            
        #end fcurve loop        
        animatedBones = list(set(animatedBones)) #remove duplicates
        staticBones = [b.name for b in bpy.data.objects[armature].pose.bones if b.name not in animatedBones]
        
        logger.info("static bones ")
        logger.info(staticBones)
        for staticBone in staticBones:
            boneData = {}
            location = bpy.data.objects[armature].pose.bones[staticBone].location
            location = [round(x, decimals) for x in location]
            scale = bpy.data.objects[armature].pose.bones[staticBone].scale
            scale = [round(x, decimals) for x in scale]            
            rotationMode = bpy.data.objects[armature].pose.bones[staticBone].rotation_mode
            
            if rotationMode != 'QUATERNION':
                #save euler
                rotation = bpy.data.objects[armature].pose.bones[staticBone].rotation_euler
                rotation = tuple(rotation)
                rotation = [round(x, decimals) for x in rotation]
            else:
                #save quaternion rotations
                rotation = bpy.data.objects[armature].pose.bones[staticBone].rotation_quaternion
                rotation = tuple(rotation)
                rotation = [round(x, decimals) for x in rotation]
            
            boneData = dict(location=location,
                            rotationMode=rotationMode,
                            rotation=rotation,
                            scale=scale)
            
            #add additional data for bendy static bone
            if self._isBendyBone(armature, staticBone):
                bendyData = self._getStaticBendyBoneDict(armature, staticBone)
                boneData.update(bendyData) #like joining to list but for dictionaries
                
            #save the static bones
            result[armature][staticKey].setdefault('1', {})[staticBone]=boneData
        
        
        #print(result)
        
        ##write to json
        with open(filePath, 'w') as outf:
            json.dump(result, outf, indent=4)
        ##
        
        
        return True

    def fromJson(self, armature, filePath, importActionName="naImportAction"):
        """import animation from json.  currently doesnt do any resetting of bones before importing animation
        todo: when applying animation > ?set all channels of bone to defaults before importing animation.
        todo: support import animation on selected bones only
        """
        if (not os.path.isfile(filePath)) or not armature:
            logger.warning("no file path or armature found to do import of animation. doing nothing")
            return False
        
        data = {}
        with open(filePath, 'r') as infile:
            data = json.load(infile)
            
        
        armatureData = data.get(armature)
        if not armatureData:
            logger.warning("no armature data found. doing nothing")
            return False
            
        animationData = armatureData.get(self._animatedKey)
        actionName = importActionName

        if animationData:
            #make an action if it doesnt exist
            action = self._createAction(armature, actionName)
            for frame, bonesDataDict in animationData.items():
                for boneName, boneData in bonesDataDict.items():
                    for channel, channelDict in boneData.items():                        
                        for arrayIndex, primaryInfo in channelDict.items():
                            #add a fcurve if doesnt exist
                            #add a keypoint 
                            self._createKeyFramePoint(armature, action, boneName=boneName, boneData=primaryInfo, frame=float(frame))

        #import static data
        staticData = armatureData.get(self._staticKey).get("1") #first frame for static data
        
        if not staticData:
            #no static data were done
            return True
            
        for boneName, boneStaticData in staticData.items():
            self._setStaticPose(armature, boneName=boneName, boneData=boneStaticData)
            
        return True

    def _setStaticPose(self, armature, boneName=None, boneData={}):
        """sets pose for static bone.
        
        ex: boneData
        "Bone": {
                    "rotationMode": "QUATERNION",
                    "location": [
                        0.0,
                        0.0,
                        0.0
                    ],
                    "rotation": [
                        1.0,
                        0.0,
                        0.0,
                        0.0
                    ],
                    "scale": [
                        1.0,
                        1.0,
                        1.0
                    ]
                },
        """
        if (not armature) or (not boneName) or (not boneData):
            logger.warning("missing inputs for creating static pose. doing nothing")
            return False
        
        skipItems = ["rotationMode"] #things in boneData that are not channels that will be set
        for channel, value in boneData.items():
            if channel in skipItems:
                continue
            #handle rotations
            if (channel == "rotation"):
                if (boneData.get("rotationMode") == "QUATERNION"):
                    #set quaternion rotation
                    setattr(bpy.data.objects[armature].pose.bones[boneName], "rotation_quaternion", value)
                else:
                    #set euler rotation
                    setattr(bpy.data.objects[armature].pose.bones[boneName], "rotation_euler", value) 
                #were done with this channel move to next
                continue
                
            #set the pose
            setattr(bpy.data.objects[armature].pose.bones[boneName], channel, value)        
            
        
    def _createKeyFramePoint(self, armature, action, boneName=None, boneData={}, frame=None):
        """create a key frame point using given data. add to appropriate fcurve. add an fcurve if doesnt exist. add keyframe point.
        todo also modify tangents appropriately
        @param armature (str) name of armature
        @param action (str) name of action
        
        ex: boneData
                  {
                    "location": {
                        "0": {
                            "array_index": 0,
                            "handle_right_type": "AUTO_CLAMPED",
                            "handle_left": [
                                11.31506,
                                0.0
                            ],
                            "value": 0.0,
                            "handle_right": [
                                20.68494,
                                0.0
                            ],
                            "interpolation": "BEZIER",
                            "channel": "location",
                            "handle_left_type": "AUTO_CLAMPED"
                        },
                        "1": {
                            "array_index": 1,
                            "handle_right_type": "AUTO_CLAMPED",
                            "handle_left": [
                                11.31506,
                                0.0
                            ],
                            "value": 0.0,
                            "handle_right": [
                                20.68494,
                                0.0
                            ],
                            "interpolation": "BEZIER",
                            "channel": "location",
                            "handle_left_type": "AUTO_CLAMPED"
                        }
                     }
                 }
        """
        if (not armature) or (not action) or (not boneData) or (not frame):
            logger.warning("missing inputs for creating keyframe point. doing nothing")
            return False
            
        
        boneAttribute = boneData.get("channel") #todo switch out to namedtuple
        attributeIndex = boneData.get("array_index")
        #todo asserts
        
        #todo make an fcurve if it doesnt exist. if it exists use existing fcurve
        fcurveIndex = self._getFcurveIndex(armature, boneName=boneName, boneAttribute=boneAttribute, attributeIndex=attributeIndex)
        
        if fcurveIndex is None:
            logger.warning("couldnt determine an fcurve. doing nothing")
            return False
            
        #insert a keyframe point
        value = boneData.get("value")
        logger.info("number fcurve>>>")
        logger.info(len(bpy.data.objects[armature].animation_data.action.fcurves))
        logger.info("fcrv index {0} frame {1} value {2}".format(fcurveIndex, frame, value))
        
        fcurve = bpy.data.objects[armature].animation_data.action.fcurves[fcurveIndex]
        keyframePoint = fcurve.keyframe_points.insert(frame, value)
        
        #set tangents and interpolation of keyframe
        extraAttrs = ["handle_right_type",
                        "handle_left_type",
                        "handle_left",
                        "handle_right",
                        "interpolation"]
        
        for extraAttr in extraAttrs:
            val = boneData.get(extraAttr)
            if val:
                setattr(keyframePoint, extraAttr, val)
        #####
        
        
    def _createAction(self, armature, actionName):
        """
        @return (str) name of created action
        """
        if (not armature in bpy.data.objects) or (not actionName):
            logger.warning("could not find inputs for creating action. doing nothing")
            return ''
            
        if actionName not in bpy.data.actions:
            actionObj = bpy.data.actions.new(actionName)
            bpy.data.objects[armature].animation_data.action = actionObj
            
        return actionName
            
    def _getFcurveIndex(self, armature, boneName='', boneAttribute='', attributeIndex=None):
        """todo: cleans this up > for creating an fcurve
        @return (int) the index of the fcurve for this bone and attribute
        """
        #todo check inputs
        dataPath = 'pose.bones["{bone}"].{channel}'.format(bone=boneName, channel=boneAttribute)
        
        #is fcurve already created
        isCurveExist = False
        allCreatedFcurves = bpy.data.objects[armature].animation_data.action.fcurves or []

        if not allCreatedFcurves:
            fcurve = bpy.data.objects[armature].animation_data.action.fcurves.new(dataPath, attributeIndex) 
            return 0 #since no fcurves have been created use 0 index

        foundIndex = None
        for i in range(0,len(allCreatedFcurves)):
            fcurve = allCreatedFcurves[i]
            if (fcurve.data_path == dataPath) and (fcurve.array_index == attributeIndex):
                #found existing curve
                isCurveExist = True
                foundIndex = i
                break
        
        if isCurveExist: 
            return foundIndex
        
        logger.warning("testing >>>>")
        logger.warning(attributeIndex)
        logger.warning(dataPath)
        fcurve = bpy.data.objects[armature].animation_data.action.fcurves.new(dataPath, attributeIndex) 
        return len(allCreatedFcurves)-1 #assuming newest fcurve is after last index. todo confirm this
        
        
    def _getStaticBendyBoneDict(self, armature, bone):
        """returns empty dictionary if not a bendy bone
        """
        isBendyBone = self._isBendyBone(armature, bone)
        if not isBendyBone:
            return {}
        
        result = {}
        channels=[  "bbone_curveinx",
                    "bbone_curveiny",
                    "bbone_curveoutx",
                    "bbone_curveouty",
                    "bbone_rollin",
                    "bbone_rollout",
                    "bbone_scalein",
                    "bbone_scaleout"
                    ]
        
        for channel in channels:
             result[channel] = eval('bpy.data.objects["{armat}"].pose.bones["{bone}"].{chan}'.format(armat=armature, bone=bone, chan=channel))
        
        return result

    def _isBendyBone(self, armature, bone):
        #todo assert bone in armature
        if bpy.data.objects[armature].data.bones[bone].bbone_segments > 1:
            return True
            
        return False