From e8c81933f195260f2bbf4491df7ed69c38464146 Mon Sep 17 00:00:00 2001 From: Hipstercat Date: Thu, 25 Feb 2021 19:31:12 +0100 Subject: [PATCH] refactor anm --- lib/skn_anm.py | 281 ++++++++++++++++++++++--------------------------- 1 file changed, 126 insertions(+), 155 deletions(-) diff --git a/lib/skn_anm.py b/lib/skn_anm.py index 95df191..22f108e 100644 --- a/lib/skn_anm.py +++ b/lib/skn_anm.py @@ -5,10 +5,27 @@ import math import prettyprinter from prettyprinter import pprint - prettyprinter.install_extras() +def with_offset(func): + def wrapped_do_offset(*args, **kwargs): + offset = kwargs.pop("offset", None) + + b = args[1] + old = b.tell() + if offset: + b.seek(offset) + + r = func(*args, **kwargs) + + if offset: + b.seek(old) + return r + + return wrapped_do_offset + + NAM_FlagHasHandles = (1 << 0) NAM_FlagCaseInsensitive = (1 << 1) ANM_NODE_FLAG_POS_ANIMATED = (1 << 0) @@ -30,15 +47,11 @@ class Vec3: return math.sqrt(math.pow(self.x - other.x, 2) + math.pow(self.y - other.y, 2) + math.pow(self.z - other.z, 2)) @classmethod - def create(cls, b, offset=None): - old = b.tell() - if offset: - b.seek(offset) + @with_offset + def create(cls, b): x = read_float_and_print("Vec3->x", b) y = read_float_and_print("Vec3->y", b) z = read_float_and_print("Vec3->z", b) - if offset: - b.seek(old) return cls(x=x, y=y, z=z) @@ -47,16 +60,12 @@ class Vec3Indexed(Vec3): index: int @classmethod - def create(cls, b, offset=None): - old = b.tell() - if offset: - b.seek(offset) + @with_offset + def create(cls, b): x = read_float_and_print("Vec3->x", b) y = read_float_and_print("Vec3->y", b) z = read_float_and_print("Vec3->z", b) index = read_int_and_print("Vec3->index", b) - if offset: - b.seek(old) return cls(x=x, y=y, z=z, index=index) @@ -65,19 +74,56 @@ class Vec4(Vec3): w: float @classmethod - def create(cls, b, offset=None): - old = b.tell() - if offset: - b.seek(offset) + @with_offset + def create(cls, b): x = read_float_and_print("Vec4->x", b) y = read_float_and_print("Vec4->y", b) z = read_float_and_print("Vec4->z", b) w = read_float_and_print("Vec4->w", b) - if offset: - b.seek(old) return cls(x=x, y=y, z=z, w=w) +@dataclass +class AnmKey: + time: int + + @classmethod + @with_offset + def create(cls, b): + time = read_byte_and_print("ANM_Key->Time", b) + return cls(time=time) + + +@dataclass +class AnmKeys: + keys_type: int + ort_before: int + ort_after: int + flags: int + p_interpolate: int + num_keys: int + keys: list + + @classmethod + @with_offset + def create(cls, b): + keys_type = read_byte_and_print("ANM_Keys->Type", b) + ort_before = read_byte_and_print("ANM_Keys->ORTBefore", b) + ort_after = read_byte_and_print("ANM_Keys->ORTAfter", b) + flags = read_byte_and_print("ANM_Keys->Flags", b) + p_interpolate = read_int_and_print("ANM_Keys->Interpolate", b) + num_keys = read_int_and_print("ANM_Keys->NumKeys", b) + p_keys = read_int_and_print("ANM_Keys->Keys", b) + keys = [] + _old = b.tell() + b.seek(p_keys) + for _ in range(num_keys): + keys.append(AnmKey.create(b)) + b.seek(_old) + return cls(keys_type=keys_type, ort_before=ort_before, ort_after=ort_after, flags=flags, + p_interpolate=p_interpolate, num_keys=num_keys, keys=keys) + + @dataclass class AnmObj: custom: int = field(repr=False) @@ -88,10 +134,8 @@ class AnmObj: clusters: List['AnmCluster'] @classmethod - def create(cls, b, offset=None): - old = b.tell() - if offset: - b.seek(offset) + @with_offset + def create(cls, b): custom = read_int_and_print("ANM_Obj->Custom", b) skin = read_int_and_print("ANM_Obj->Skin", b) zero = read_int_and_print("ANM_Obj->ZeroCluster", b) @@ -106,9 +150,8 @@ class AnmObj: clusters.append(cluster) b.seek(_old1) - if offset: - b.seek(old) - r = cls(custom=custom, skin=skin, zero_cluster=zero, total_number_vertices=total_vertices, num_clusters=num_clusters, clusters=clusters) + r = cls(custom=custom, skin=skin, zero_cluster=zero, total_number_vertices=total_vertices, + num_clusters=num_clusters, clusters=clusters) for cluster in clusters: cluster.obj = r return r @@ -125,15 +168,13 @@ class AnmCluster: bone_name: str = "" @classmethod - def create(cls, b, offset=None): - old = b.tell() - if offset: - b.seek(offset) + @with_offset + def create(cls, b): custom = read_int_and_print("ANM_Cluster->Custom", b) obj_offset = read_int_and_print("ANM_Cluster->Obj", b) if obj_offset == 0: # ugly hack, dunno why we need that - _ = read_int_and_print("ANM_Cluster->Custom", b) - _ = read_int_and_print("ANM_Cluster->Obj", b) + custom = read_int_and_print("ANM_Cluster->Custom", b) + obj_offset = read_int_and_print("ANM_Cluster->Obj", b) obj = None # is set by parent Obj handle = read_int_and_print("ANM_Cluster->Handle", b) bbox1 = Vec3.create(b) @@ -145,24 +186,21 @@ class AnmCluster: v = Vec3Indexed.create(b) vertices.append(v) - if offset: - b.seek(old) - return cls(custom=custom, obj=obj, handle=handle, bounding_box=bounding_box, num_vertices=num_vertices, vertices=vertices) + return cls(custom=custom, obj=obj, handle=handle, bounding_box=bounding_box, num_vertices=num_vertices, + vertices=vertices) class NamDictionnary: @staticmethod - def create(b, cls, depth=1, offset=None): + @with_offset + def create(cls, b, depth=1): r = {} - old = b.tell() - if offset: - b.seek(offset) - read_int_and_print("NAM_Dictionnary->pMemGroup", b) namemapper = read_int_and_print("NAM_Dictionnary->pNameMapper", b) flags = read_int_and_print("NAM_Dictionnary->Flags", b) - print("Flags: NAM_FlagHasHandles: %s, NAM_FlagCaseInsensitive: %s" % (flags & NAM_FlagHasHandles, flags & NAM_FlagCaseInsensitive)) + print("Flags: NAM_FlagHasHandles: %s, NAM_FlagCaseInsensitive: %s" % ( + flags & NAM_FlagHasHandles, flags & NAM_FlagCaseInsensitive)) read_short_and_print("NAM_Dictionnary->NameHandle", b) b.seek(namemapper) @@ -177,9 +215,7 @@ class NamDictionnary: read_int_and_print("TotalKeys", b) p_key_contexts = read_int_and_print("pKeyContexts", b) - print("MAP_Image->ANM_Skin->NAM_Dictionnary->MAP_Mapper->MAP_KeyContext") b.seek(p_key_contexts) - print("\tMAP_Image->ANM_Skin->NAM_Dictionnary->MAP_Mapper->MAP_KeyContext->MAP_BinSearchContext") _ = read_int_and_print("\tSize", b) read_int_and_print("\tpCallback", b) read_int_and_print("\tIndex", b) @@ -193,7 +229,6 @@ class NamDictionnary: pindx = read_int(b) for _ in range(i): pindx = read_int(b) - print("LUT INDEX: %s" % pindx) b.seek(pindx) if not flags & NAM_FlagHasHandles: read_short_and_print("pPtr", b) @@ -201,9 +236,9 @@ class NamDictionnary: val = read_int_and_print("Value", b) name = read_string_until_none_and_print("name", b) if depth > 1: - item = NamDictionnary.create(b, cls, depth-1, val) + item = NamDictionnary.create(cls, b, depth - 1, offset=val) else: - item = cls.create(b, val) + item = cls.create(b, offset=val) r[name] = item else: refcnt = read_short_and_print("RefCnt", b) @@ -213,9 +248,6 @@ class NamDictionnary: bone_name = read_string_until_none_and_print("BoneName", b) r[str(i)] = {"BoneName": bone_name} - if offset: - b.seek(old) - return r @@ -224,22 +256,13 @@ class GeoMat4x4: m: List[float] @classmethod - def create(cls, b, offset=None): - old = b.tell() - if offset: - b.seek(offset) - + @with_offset + def create(cls, b): m = [] for _ in range(4): for _j in range(4): m.append(read_float(b)) - - r = cls(m=m) - - if offset: - b.seek(old) - - return r + return cls(m=m) @dataclass @@ -248,20 +271,11 @@ class GeoMat: s: GeoMat4x4 @classmethod - def create(cls, b, offset=None): - old = b.tell() - if offset: - b.seek(offset) - + @with_offset + def create(cls, b): mat_class = read_int(b) s = GeoMat4x4.create(b) - - r = cls(mat_class=mat_class, s=s) - - if offset: - b.seek(old) - - return r + return cls(mat_class=mat_class, s=s) @dataclass @@ -272,22 +286,13 @@ class GeoAffine: affine_class: int @classmethod - def create(cls, b, offset=None): - old = b.tell() - if offset: - b.seek(offset) - + @with_offset + def create(cls, b): trans = Vec3.create(b) quat = Vec4.create(b) scale = Vec3.create(b) affine_class = read_int(b) - - r = cls(transform=trans, quat=quat, scale=scale, affine_class=affine_class) - - if offset: - b.seek(old) - - return r + return cls(transform=trans, quat=quat, scale=scale, affine_class=affine_class) @dataclass @@ -296,20 +301,11 @@ class AnmTransform: affine: GeoAffine @classmethod - def create(cls, b, offset=None): - old = b.tell() - if offset: - b.seek(offset) - + @with_offset + def create(cls, b): time = read_int(b) affine = GeoAffine.create(b) - - r = cls(time=time, affine=affine) - - if offset: - b.seek(old) - - return r + return cls(time=time, affine=affine) @dataclass @@ -326,16 +322,13 @@ class AnmNode: stamp: int geo_mat: GeoMat transform: AnmTransform - pos_keys: int - rot_keys: int - scl_keys: int + pos_keys: Union[AnmKeys, int] + rot_keys: Union[AnmKeys, int] + scl_keys: Union[AnmKeys, int] @classmethod - def create(cls, b, offset=None): - old = b.tell() - if offset: - b.seek(offset) - + @with_offset + def create(cls, b): custom = read_int_and_print("ANM_Node->Custom", b) node_type = read_short_and_print("ANM_Node->Type", b) x_form_type = read_short_and_print("ANM_Node->XformType", b) @@ -352,27 +345,19 @@ class AnmNode: rot_keys = read_int_and_print("ANM_Node->pRotKeys", b) scl_keys = read_int_and_print("ANM_Node->pSclKeys", b) - """ - if nodeflags & ANM_NODE_FLAG_POS_ANIMATED: - b.seek(posKeys) - read_byte_and_print("\t\tType", b) - read_byte_and_print("\t\tORTBefore", b) - read_byte_and_print("\t\tORTAfter", b) - read_byte_and_print("\t\tFlags", b) - read_int_and_print("\t\tpInterpolate", b) - read_int_and_print("\t\tNumberOfKeys", b) - read_int_and_print("\t\tpKeys", b) - """ + if pos_keys > 0: + pos_keys = AnmKeys.create(b, offset=pos_keys) - r = cls(custom=custom, node_type=node_type, x_form_type=x_form_type, flags=flags, p_anim=p_anim, - p_parent=p_parent, p_target=p_target, p_usercall=p_usercall, p_usercontext=p_usercontext, - stamp=stamp, geo_mat=geo_mat, transform=transform, pos_keys=pos_keys, rot_keys=rot_keys, - scl_keys=scl_keys) + if rot_keys > 0: + rot_keys = AnmKeys.create(b, offset=rot_keys) - if offset: - b.seek(old) + if scl_keys > 0: + scl_keys = AnmKeys.create(b, offset=scl_keys) - return r + return cls(custom=custom, node_type=node_type, x_form_type=x_form_type, flags=flags, p_anim=p_anim, + p_parent=p_parent, p_target=p_target, p_usercall=p_usercall, p_usercontext=p_usercontext, + stamp=stamp, geo_mat=geo_mat, transform=transform, pos_keys=pos_keys, rot_keys=rot_keys, + scl_keys=scl_keys) @dataclass @@ -403,11 +388,8 @@ class AnmAnim: name: str @classmethod - def create(cls, b, offset=None): - old = b.tell() - if offset: - b.seek(offset) - + @with_offset + def create(cls, b): custom = read_int_and_print("ANM_Anim->Custom", b) version = read_int_and_print("ANM_Anim->FileVersion", b) image = read_int_and_print("ANM_Anim->pImage", b) @@ -430,19 +412,14 @@ class AnmAnim: otheranim = read_int_and_print("ANM_AnimExtra->OtherAnim", b) name = read_string_until_none(b) - treedict = NamDictionnary.create(b, AnmNode, depth=2, offset=treedict) - defaultobjdict = NamDictionnary.create(b, AnmNode, offset=defaultobjdict) + treedict = NamDictionnary.create(AnmNode, b, depth=2, offset=treedict) + defaultobjdict = NamDictionnary.create(AnmNode, b, offset=defaultobjdict) - r = cls(custom=custom, fileversion=version, image=image, memgroup=memgroup, start_time=start_time, - end_time=end_time, ticks_per_frame=ticks_per_frame, framerate=framerate, tree_dictionnary=treedict, - default_obj_dictionnary=defaultobjdict, num_nodes=num_nodes, extra_stamp=stamp, extra_flags=flags, - extra_time1=time1, extra_stamp1=stamp1, extra_time2=time2, extra_stamp2=stamp2, extra_factor=factor, - extra_other_anim=otheranim, name=name) - - if offset: - b.seek(old) - - return r + return cls(custom=custom, fileversion=version, image=image, memgroup=memgroup, start_time=start_time, + end_time=end_time, ticks_per_frame=ticks_per_frame, framerate=framerate, tree_dictionnary=treedict, + default_obj_dictionnary=defaultobjdict, num_nodes=num_nodes, extra_stamp=stamp, extra_flags=flags, + extra_time1=time1, extra_stamp1=stamp1, extra_time2=time2, extra_stamp2=stamp2, extra_factor=factor, + extra_other_anim=otheranim, name=name) @dataclass @@ -468,11 +445,8 @@ class AnmSkin: cluster.bone_name = nodedict[k]["BoneName"] @classmethod - def create(cls, b, offset=None): - old = b.tell() - if offset: - b.seek(offset) - + @with_offset + def create(cls, b): custom = read_int_and_print("pCustom", b) version = read_int_and_print("FileVersion", b) image = read_int_and_print("pImage", b) @@ -485,17 +459,14 @@ class AnmSkin: num_clusters = read_int_and_print("TotalNumberOfClusters", b) name = read_string_until_none_and_print("Name", b) - nodedict = NamDictionnary.create(b, AnmObj, offset=node) - treedict = NamDictionnary.create(b, AnmObj, depth=2, offset=tree) - defaultobjdict = NamDictionnary.create(b, AnmObj, offset=defaultobj) + nodedict = NamDictionnary.create(AnmObj, b, offset=node) + treedict = NamDictionnary.create(AnmObj, b, depth=2, offset=tree) + defaultobjdict = NamDictionnary.create(AnmObj, b, offset=defaultobj) # AnmSkin.link(nodedict, defaultobjdict) - r = cls(custom=custom, file_version=version, image=image, memgroup=memgroup, node_dictionnary=nodedict, tree_dictionnary=treedict, default_obj_dictionnary=defaultobjdict, total_number_trees=num_trees, total_number_objs=num_objs, total_number_clusters=num_clusters, name=name) - - if offset: - b.seek(old) - - return r + return cls(custom=custom, file_version=version, image=image, memgroup=memgroup, node_dictionnary=nodedict, + tree_dictionnary=treedict, default_obj_dictionnary=defaultobjdict, total_number_trees=num_trees, + total_number_objs=num_objs, total_number_clusters=num_clusters, name=name) @dataclass @@ -514,4 +485,4 @@ class AnmFile: read_int_and_print("TotalSize", b) read_int_and_print("TotalBlocks", b) base = read_int_and_print("pBase", b) - return cls.create(b, base) + return cls.create(b, offset=base)