blender import psa

bipsa.py

bl_info = {
    "name": "Import Unreal Skeleton Anim (.psa)",
    "author": "Darknet",
    "version": (2, 0),
    "blender": (2, 6, 2),
    "location": "File > Import > Skeleton Anim (.psa)",
    "description": "Import Skeleleton Anim",
    "warning": "",
    "wiki_url": "http://wiki.blender.org/index.php/Extensions:2.5/Py/"
        "Scripts/Import-Export/Unreal_psk_psa",
    "tracker_url": "https://projects.blender.org/tracker/index.php?"\
        "func=detail&aid=21366",
    "category": "Import-Export"}
    
import bpy
import mathutils
import math
from mathutils import *
from string import *
from struct import *
from math import *
from bpy.props import *

class psa_bone:
    name=""
    Transform=None
    parent=None
    def __init__(self):
        self.name=""
        self.Transform=None
        self.parent=None


def psaimport(filename,context):
    print ("--------------------------------------------------")
    print ("---------SCRIPT EXECUTING PYTHON IMPORTER---------")
    print ("--------------------------------------------------")
    print ("Importing file: ", filename)
    psafile = open(filename,'rb')
    debug = True
    if (debug):
        logpath = filename.replace(".psa", ".txt")
        print("logpath:",logpath)
        logf = open(logpath,'w')
    def printlog(strdata):
        if (debug):
            logf.write(strdata)
    def printlogplus(name,data):
        if (debug):
            logf.write(str(name)+'\n')
            if isinstance(data,bytes):
                logf.write(str(bytes.decode(data).strip(bytes.decode(b'\x00'))))
            else:
                logf.write(str(data))
            logf.write('\n')
            
    printlog('-----------Log File------------\n')
    #General Header
    indata = unpack('20s3i',psafile.read(32))
    printlogplus('ChunkID',indata[0])
    printlogplus('TypeFlag',indata[1])
    printlogplus('DataSize',indata[2])
    printlogplus('DataCount',indata[3])
    #Bones Header
    indata = unpack('20s3i',psafile.read(32))
    printlogplus('ChunkID',indata[0])
    printlogplus('TypeFlag',indata[1])
    printlogplus('DataSize',indata[2])
    printlogplus('DataCount',indata[3])
    #Bones Data
    BoneIndex2NamePairMap = {}
    BoneNotFoundList = []
    printlog("Name|Flgs|NumChld|PrntIdx|Qx|Qy|Qz|Qw|LocX|LocY|LocZ|Length|XSize|YSize|ZSize\n")
    recCount = indata[3]
    counter = 0
    nobonematch = True
    while counter < recCount:
        indata = unpack('64s3i11f',psafile.read(120))
        #printlogplus('bone',indata[0])
        bonename = str(bytes.decode(indata[0]).strip(bytes.decode(b'\x00')))
        if bonename in bpy.data.armatures['armaturedata'].bones.keys():
            BoneIndex2NamePairMap[counter] = bonename
            print('find bone',bonename)
            nobonematch = False
        else:
            print('can not find the bone:',bonename)
            BoneNotFoundList.append(counter)

        counter += 1
    
    if nobonematch:
        print('no bone was match so skip import!')
        return
        
    #Animations Header
    indata = unpack('20s3i',psafile.read(32))
    printlogplus('ChunkID',indata[0])
    printlogplus('TypeFlag',indata[1])
    printlogplus('DataSize',indata[2])
    printlogplus('DataCount',indata[3])
    #Animations Data
    recCount = indata[3]
    counter = 0
    Raw_Key_Nums = 0
    Action_List = []
    while counter < recCount:
        indata = unpack('64s64s4i3f3i',psafile.read(64+64+4*4+3*4+3*4))
        printlogplus('Name',indata[0])
        printlogplus('Group',indata[1])
        printlogplus('totalbones',indata[2])
        printlogplus('NumRawFrames',indata[-1])
        Name = str(bytes.decode(indata[0]).strip(bytes.decode(b'\x00')))
        Group = str(bytes.decode(indata[1]).strip(bytes.decode(b'\x00')))
        totalbones = indata[2]
        NumRawFrames = indata[-1]
        
        Raw_Key_Nums += indata[2] * indata[-1]
        Action_List.append((Name,Group,totalbones,NumRawFrames))
        
        counter += 1
        
    
    #Raw keys Header
    Raw_Key_List = []
    indata = unpack('20s3i',psafile.read(32))
    printlogplus('ChunkID',indata[0])
    printlogplus('TypeFlag',indata[1])
    printlogplus('DataSize',indata[2])
    printlogplus('DataCount',indata[3])
    if(Raw_Key_Nums != indata[3]):
        print('error! Raw_Key_Nums Inconsistent')
        return
    #Raw keys Data
    recCount = Raw_Key_Nums
    counter = 0
    while counter < recCount:
        indata = unpack('3f4f1f',psafile.read(3*4+4*4+4))
        pos = mathutils.Vector((indata[0], indata[1], indata[2]))
        quat = mathutils.Quaternion((indata[6], indata[3], indata[4], indata[5]))
        time = indata[7]
        Raw_Key_List.append((pos,quat,time))
        counter += 1
    #Scale keys Header,Scale keys Data,Curve keys Header,Curve keys Data
    curFilePos = psafile.tell()
    psafile.seek(0,2)
    endFilePos = psafile.tell()
    if curFilePos == endFilePos:
        print('no Scale keys,Curve keys')
     
     
    #build the animation line
    if bpy.ops.object.mode_set.poll():
        bpy.ops.object.mode_set(mode='OBJECT', toggle=False)
        
    
    NeededBoneMatrix = {}
    for bone in bpy.data.armatures['armaturedata'].bones:
        name = bone.name
        ori_matrix = bone.matrix
        matrix = bone.matrix_local.to_3x3()
        bone_rest_matrix = Matrix(matrix)
        
        
        #bone_rest_matrix = bone.matrix_local.to_3x3()
        #bone_rest_matrix = bone.matrix_local.to_quaternion().conjugated().to_matrix()

        bone_rest_matrix_inv = Matrix(bone_rest_matrix)
        bone_rest_matrix_inv.invert()

        bone_rest_matrix_inv.resize_4x4()
        bone_rest_matrix.resize_4x4()
        NeededBoneMatrix[name] = (bone_rest_matrix,bone_rest_matrix_inv,ori_matrix)
        

    #build tmp pose bone tree
    psa_bones = {}
    for bone in bpy.data.objects['ArmObject'].pose.bones:
        _psa_bone = psa_bone()
        _psa_bone.name = bone.name
        _psa_bone.Transform = bone.matrix
        if bone.parent != None:
            _psa_bone.parent = psa_bones[bone.parent.name]
        else:
            _psa_bone.parent = None
        psa_bones[bone.name] = _psa_bone
    
    raw_key_index = 0



    for raw_action in Action_List:
        Name            = raw_action[0]
        Group           = raw_action[1]
        Totalbones      = raw_action[2]
        NumRawFrames    = raw_action[3]
        context.scene.update()
        object = bpy.data.objects['ArmObject']
        object.animation_data_create()
        action = bpy.data.actions.new(name=Name)    
        object.animation_data.action = action
        for i in range(NumRawFrames):
            context.scene.frame_set(i+1)
            pose_bones = object.pose.bones
            for j in range(Totalbones):
                if j not in BoneNotFoundList:
                    bName = BoneIndex2NamePairMap[j]
                    pbone           = psa_bones[bName]
                    pos = Raw_Key_List[raw_key_index][0]
                    quat = Raw_Key_List[raw_key_index][1]
           
                    mat = Matrix()
                    if pbone.parent != None:
                        quat = quat.conjugated()
                        mat = Matrix.Translation(pos) * quat.to_matrix().to_4x4()
                        mat = pose_bones[bName].parent.matrix * mat 
                        #mat = pbone.parent.Transform * mat
                    else:
                        mat = pbone.Transform * Matrix.Translation(pos) * quat.to_matrix().to_4x4()

                    pose_bones[bName].matrix = mat
                    pbone.Transform = mat

                raw_key_index += 1
            
            

            #bpy.data.meshes[1]
            for bone in pose_bones:              
                bone.matrix = psa_bones[bone.name].Transform
                bone.keyframe_insert("rotation_quaternion")
                bone.keyframe_insert("location")
            

            def whirlSingleBone(pose_bone,quat):
                bpy.context.scene.update()
                #record child's matrix and origin rotate
                hymat = Quaternion((0.707,-0.707,0,0)).inverted().to_matrix().to_4x4()
                
                
                children_infos = {}
                childrens = pose_bone.children
                for child in childrens:
                    armmat = bpy.data.armatures['armaturedata'].bones[child.name].matrix.copy().to_4x4()                    
                    cmat = child.matrix.copy() * armmat.inverted() * hymat.inverted()
                    pos = cmat.to_translation()
                    rotmat = cmat.to_3x3()
                    children_infos[child] = (armmat,pos,rotmat)
                
                #whirl this bone by quat
                pose_bone.matrix *= quat.to_matrix().to_4x4()
                pose_bone.keyframe_insert("location")
                pose_bone.keyframe_insert("rotation_quaternion")
                bpy.context.scene.update()
                #set back children bon to original position 
                #reverse whirl child bone by quat.inverse()
                
                for child in childrens:
                    armmat = children_infos[child][0]
                    pos = children_infos[child][1]
                    rotmat = children_infos[child][2] 

                    child.matrix = Matrix.Translation(pos) * rotmat.to_4x4() * hymat * armmat 
                    child.keyframe_insert("location") 
                    child.keyframe_insert("rotation_quaternion")
                
                


            for bone in pose_bones:
                if bone.parent != None:
                    whirlSingleBone(bone,Quaternion((0.707,0,0,-0.707)))
                else:
                    bone.rotation_quaternion *= Quaternion((0.707,-0.707,0,0))*Quaternion((0.707,0,0,-0.707))
                    bone.keyframe_insert("rotation_quaternion")
            
        break
        
    context.scene.frame_set(0)
    if(debug):
        logf.close()
    
def getInputFilename(self,filename,context):
    checktype = filename.split('\\')[-1].split('.')[1]
    if checktype.lower() != 'psa':
        print ("  Selected file = ",filename)
        raise (IOError, "The selected input file is not a *.psa file")
        #self.report({'INFO'}, ("Selected file:"+ filename))
    else:
        psaimport(filename,context)

class IMPORT_OT_psa(bpy.types.Operator):
    '''Load a skeleton anim psa File'''
    bl_idname = "import_scene.psa"
    bl_label = "Import PSA"
    bl_space_type = "PROPERTIES"
    bl_region_type = "WINDOW"

    filepath = StringProperty(
            subtype='FILE_PATH',
            )
    filter_glob = StringProperty(
            default="*.psa",
            options={'HIDDEN'},
            )

    def execute(self, context):
        getInputFilename(self,self.filepath,context)
        return {'FINISHED'}

    def invoke(self, context, event):
        wm = context.window_manager
        wm.fileselect_add(self)
        return {'RUNNING_MODAL'}  

def menu_func(self, context):
    self.layout.operator(IMPORT_OT_psa.bl_idname, text="Skeleton Anim (.psa)")

def register():
    bpy.utils.register_module(__name__)
    bpy.types.INFO_MT_file_import.append(menu_func)
    
def unregister():
    bpy.utils.unregister_module(__name__)
    bpy.types.INFO_MT_file_import.remove(menu_func)

if __name__ == "__main__":
    register()