WIP: Import bones

Signed-off-by: Lucas Schwiderski <lucas@lschwiderski.de>

addons/bitsquid/import_bsi.py

@@ -178,12 +178,12 @@ def find(arr, f):
             return val, key
 
 
-def create_object(self, context, name, node_data, geometries):
+def create_mesh(self, context, name, node_data, geo_data):
     """
     Create a Blender object from a BSI node definition
     and additional data from the file.
     """
-    print("[create_object]", name, node_data)
+
     # A list of vectors that represent vertex locations.
     vertices = []
     # A list of vectors, where each vector contains three indices into `vertices`.
@@ -192,10 +192,8 @@ def create_object(self, context, name, node_data, geometries):
     uv_name = "UVMap"
     uvs = []
 
-    tri_count = int(geometries["indices"]["size"]) / 3
-
-    for i, index_stream in enumerate(geometries["indices"]["streams"]):
-        data_stream = geometries["streams"][i]
+    for i, index_stream in enumerate(geo_data["indices"]["streams"]):
+        data_stream = geo_data["streams"][i]
         stride = data_stream["stride"] / 4
 
         for channel in data_stream["channels"]:
@@ -230,7 +228,15 @@ def create_object(self, context, name, node_data, geometries):
                 # Blender is able to create normals from the face definition
                 # (i.e. the order in which the faces vertices were defined)
                 # and that seems to be good enough
-                self.report({'INFO'}, "Ignoring custom normals")
+                # self.report({'INFO'}, "Ignoring custom normals")
+                continue
+            elif channel["name"] in {'ALPHA', 'COLOR'}:
+                # Not sure if this can be intended as a primitive material,
+                # but I'll assume it's not. And while Blender does feature
+                # the concept of viewport-only vertex colors, that's rather
+                # low priority to implement.
+                # self.report({'INFO'}, "Ignoring vertex color data")
+                continue
             elif channel["name"] == 'TEXCOORD':
                 uv_name = "UVMap{}".format(channel["index"] + 1)
                 uv_data = data_stream["data"]
@@ -238,16 +244,16 @@ def create_object(self, context, name, node_data, geometries):
                 uvs = [uv_defs[index_stream[j]] for j in range(0, len(index_stream))]
             else:
                 # TODO: Implement other channel types
-                self.report(
-                    {'WARNING'},
-                    "Unknown channel type: {}".format(channel["name"])
-                )
+                # self.report(
+                #     {'WARNING'},
+                #     "Unknown channel type: {}".format(channel["name"])
+                # )
+                continue
 
     mesh = bpy.data.meshes.new(name)
     mesh.from_pydata(vertices, [], faces)
 
     if len(uvs) > 0:
-        print("{} has UVs".format(name))
         uv_layer = mesh.uv_layers.new(name=uv_name)
         uv_layer.data.foreach_set("uv", unpack_list(uvs))
 
@@ -271,22 +277,193 @@ def matrix_from_list(list):
     return Matrix(rows)
 
 
-def import_node(self, context, name, node_data, global_data):
-    """Import a BSI node. Recurses into child nodes."""
-    print("[import_node]", name, node_data)
-    if "geometries" in node_data:
-        geometry_name = node_data["geometries"][0]
-        print("[import_node] Building geometry '{}' for object '{}'".format(geometry_name, name))
-        geometries = global_data["geometries"][geometry_name]
-        mesh = create_object(self, context, name, node_data, geometries)
-        obj = bpy.data.objects.new(mesh.name, mesh)
-    else:
-        print("[import_node] Adding empty for '{}'".format(name))
-        obj = bpy.data.objects.new(name, None)
-        # Decrease axis size to prevent overcrowding in the viewport
-        obj.empty_display_size = 0.1
+def import_joint(
+            self,
+            context,
+            name,
+            node_data,
+            armature,
+            parent_bone,
+            parent_rotation,
+            global_data
+        ):
+    """
+    Imports a joint and all of its children.
+    In BSI (smilar to Maya) skeletons are defined as a series of joints,
+    with bones added virtually in between, when needed.
 
+    In Blender, skeletons are defined with bones, where each bone has a `head`
+    and `tail`. So we need to convert the list of joints to a series of bones
+    instead.
+    The challenge here is the fact that joints have rotation data, whereas
+    `head` and `tail` for bones don't. This changes how the position data has
+    to be treated, as it is relative to the respective previous joint.
+    Compared to the code that imports mesh objects, we can't just apply
+    the matrix to the bone and have it position itself relative to its parent.
+    Instead, we have to manually keep track of the parent's rotation.
+    """
+    if "local" not in node_data:
+        raise RuntimeError("No position value for joint '{}'".format(name))
+
+    mat = matrix_from_list(node_data["local"])
+    translation, rotation, _ = mat.decompose()
+
+    if name.endswith("_scale"):
+        # print("Skipping joint '{}'".format(name))
+        bone = parent_bone
+    else:
+        bone = armature.edit_bones.new(name)
+        if parent_bone:
+            # The values for `bone.head` and `bone.tail` are relative to their
+            # parent, so we need to apply that first.
+            bone.parent = parent_bone
+            bone.use_connect = True
+            # bone.head = parent_bone.tail
+        else:
+            bone.head = Vector((0, 0, 0))
+
+        if parent_rotation:
+            print("[import_joint] {} Parent @ Local:".format(name), parent_rotation, parent_rotation @ translation)
+            bone.tail = bone.head + (parent_rotation @ translation)
+        else:
+            bone.tail = bone.head + translation
+
+        print("[import_joint] {} Local:".format(name), translation)
+        print("[import_joint] {} Bone:".format(name), bone.head, bone.tail)
+
+    if "children" in node_data:
+        for child_name, child_data in node_data["children"].items():
+            if child_data["parent"] != name:
+                raise RuntimeError(
+                    "Assigned parent '{}' doesn't match actual parent node '{}'"
+                    .format(child_data["parent"], name)
+                )
+
+            if child_name.startswith("j_"):
+                child_bone = import_joint(
+                    self,
+                    context,
+                    child_name,
+                    child_data,
+                    armature,
+                    bone,
+                    rotation,
+                    global_data
+                )
+                child_bone.parent = bone
+            else:
+                # DEBUG: ignoring these for now
+                continue
+                # Not entirely sure, but I think these are considered
+                # "controller nodes" in Maya. Would make sense based on
+                # name.
+                if "children" in child_data:
+                    raise RuntimeError(
+                        "Controller node '{}' has children."
+                        .format(child_name)
+                    )
+
+                if "geometries" not in child_data:
+                    raise RuntimeError(
+                        "Controller node '{}' has no geometry."
+                        .format(child_name)
+                    )
+
+                child_obj = import_node(
+                    self,
+                    context,
+                    child_name,
+                    child_data,
+                    global_data,
+                )
+                # TODO: How to parent to a bone?
+                child_obj.parent = bone
+
+    return bone
+
+
+def import_armature(self, context, name, node_data, global_data):
+    armature = context.blend_data.armatures.new(name)
+
+    # An armature itself cannot exist in the view layer.
+    # We need to create an object from it
+    obj = bpy.data.objects.new(name, armature)
+    context.collection.objects.link(obj)
+
+    # Armature needs to be in EDIT mode to allow adding bones
+    context.view_layer.objects.active = obj
+    bpy.ops.object.mode_set(mode='EDIT', toggle=False)
+
+    if "local" not in node_data:
+        raise RuntimeError("No position value for joint '{}'".format(name))
+
+    mat = matrix_from_list(node_data["local"])
+    obj.matrix_local = mat
+    # DEBUG
+    mat_loc, mat_rot, mat_scale = mat.decompose()
+    print("[import_joint] {}".format(name), mat_loc, mat_rot)
+
+    if "children" in node_data:
+        for child_name, child_data in node_data["children"].items():
+            if child_data["parent"] != name:
+                raise RuntimeError(
+                    "Assigned parent '{}' doesn't match actual parent node '{}'"
+                    .format(child_data["parent"], name)
+                )
+
+            if child_name.startswith("j_"):
+                import_joint(
+                    self,
+                    context,
+                    child_name,
+                    child_data,
+                    armature,
+                    None,
+                    None,
+                    global_data
+                )
+            else:
+                # DEBUG: Ignoring these for now
+                continue
+                # Not entirely sure, but I think these are considered
+                # "controller nodes" in Maya. Would make sense based on
+                # name.
+                if "children" in child_data:
+                    raise RuntimeError(
+                        "Controller node '{}' has children."
+                        .format(child_name)
+                    )
+
+                child_obj = import_node(
+                    self,
+                    context,
+                    child_name,
+                    child_data,
+                    global_data,
+                )
+                child_obj.parent = obj
+
+    # Disable EDIT mode
+    context.view_layer.objects.active = obj
+    bpy.ops.object.mode_set(mode='OBJECT')
+
+    return obj
+
+
+def import_geometry(self, context, name, node_data, global_data):
+    if len(node_data["geometries"]) > 1:
+        self.report(
+            {'WARNING'},
+            "More than one geometry for node '{}'.".format(name)
+        )
+
+    geometry_name = node_data["geometries"][0]
+    geo_data = global_data["geometries"][geometry_name]
+
+    mesh = create_mesh(self, context, name, node_data, geo_data)
+    obj = bpy.data.objects.new(mesh.name, mesh)
     obj.matrix_world = Matrix()
+
     if "local" in node_data:
         mat = matrix_from_list(node_data["local"])
         obj.matrix_local = mat
@@ -295,8 +472,64 @@ def import_node(self, context, name, node_data, global_data):
         for child_name, child_data in node_data["children"].items():
             if child_data["parent"] != name:
                 raise RuntimeError(
-                    "Assigned parent '%s' doesn't match actual parent node '%s'"
-                    % (child_data["parent"], name)
+                    "Assigned parent '{}' doesn't match actual parent node '{}'"
+                    .format(child_data["parent"], name)
+                )
+
+            child_obj = import_node(
+                self,
+                context,
+                child_name,
+                child_data,
+                global_data,
+            )
+
+            if not isinstance(child_obj, bpy.types.Object):
+                raise RuntimeError(
+                    "Node of type '{}' cannot be child of a geometry node."
+                    .format(type(child_obj))
+                )
+
+            child_obj.parent = obj
+
+    # Make sure all objects are linked to the current collection.
+    # Otherwise they won't show up in the outliner.
+    collection = context.collection
+    collection.objects.link(obj)
+    return obj
+
+
+def import_node(self, context, name, node_data, global_data):
+    """Import a BSI node. Recurses into child nodes."""
+    has_geometry = "geometries" in node_data
+    is_joint = name in global_data["joint_index"]
+
+    if is_joint:
+        if has_geometry:
+            raise RuntimeError("Found geometry data in joint '{}'".format(name))
+
+        if not name.startswith("j_"):
+            raise RuntimeError("Invalid name for joint: '{}".format(name))
+
+        self.report({'INFO'}, "Creating Armature '{}'".format(name))
+        return import_armature(self, context, name, node_data, global_data)
+
+    if has_geometry:
+        return import_geometry(self, context, name, node_data, global_data)
+    else:
+        if name != "root_point":
+            self.report({'WARNING'}, "Unknown kind of node: '{}'. Falling back to Empty.".format(name))
+
+        obj = bpy.data.objects.new(name, None)
+        # Decrease axis size to prevent overcrowding in the viewport
+        obj.empty_display_size = 0.1
+
+        if "children" in node_data:
+            for child_name, child_data in node_data["children"].items():
+                if child_data["parent"] != name:
+                    raise RuntimeError(
+                        "Assigned parent '{}' doesn't match actual parent node '{}'"
+                        .format(child_data["parent"], name)
                     )
 
                 child_obj = import_node(
@@ -312,27 +545,34 @@ def import_node(self, context, name, node_data, global_data):
         # Otherwise they won't show up in the outliner.
         collection = context.collection
         collection.objects.link(obj)
+
         return obj
 
 
 def load(self, context, filepath, *, relpath=None):
     global_data = parse_sjson(filepath)
 
-    print(global_data)
-
     # Nothing to do if there are no nodes
     if "nodes" not in global_data:
         self.report({'WARNING'}, "No nodes to import in {}".format(filepath))
         return {'CANCELLED'}
 
-    view_layer = context.view_layer
+    # Build joint index
+    joint_index = []
+    if "skins" in global_data:
+        for v in global_data["skins"].values():
+            for joint in v["joints"]:
+                name = joint["name"]
+                if name in global_data["geometries"]:
+                    self.report({'ERROR'}, "Found joint with mesh data.")
+                    return {'CANCELLED'}
+
+                if name not in joint_index:
+                    joint_index.append(joint['name'])
+
+    global_data["joint_index"] = joint_index
 
     for name, node_data in global_data["nodes"].items():
-        obj = import_node(self, context, name, node_data, global_data)
-
-        obj.select_set(True)
-        # One of the objects should be set as active.
-        # This is the easiest to implement and perfectly fine.
-        view_layer.objects.active = obj
+        import_node(self, context, name, node_data, global_data)
 
     return {'FINISHED'}