character_refine.py

import os
import sys
from glob import glob

import bpy
import numpy as np
from tqdm import tqdm


class HiddenPrints:
    def __init__(self, disable=False):
        self.disable = disable

    def __enter__(self):
        if self.disable:
            return
        self._original_stdout = sys.stdout
        sys.stdout = open(os.devnull, "w")

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.disable:
            return
        sys.stdout.close()
        sys.stdout = self._original_stdout


def remove_all(delete_actions=True):
    for obj in bpy.data.objects.values():
        bpy.data.objects.remove(obj, do_unlink=True)
    bpy.ops.outliner.orphans_purge(do_recursive=True)
    if delete_actions:
        for action in bpy.data.actions:
            bpy.data.actions.remove(action, do_unlink=True)


def load_file(filepath: str):
    old_objs = set(bpy.context.scene.objects)
    if filepath.endswith(".glb"):
        bpy.ops.import_scene.gltf(filepath=filepath)
    elif filepath.endswith(".fbx"):
        bpy.ops.import_scene.fbx(filepath=filepath)
    else:
        raise RuntimeError(f"Invalid input file: {filepath}")
    imported_objs = set(bpy.context.scene.objects) - old_objs
    print("Imported:", imported_objs)
    return imported_objs


def get_type_objs(obj_list=None, type="MESH"):
    if not obj_list:
        obj_list = bpy.context.scene.objects
    type_obj_list = []
    for obj in obj_list:
        if obj.type == type:
            type_obj_list.append(obj)
    return type_obj_list


def get_all_mesh_obj(obj_list=None):
    return get_type_objs(obj_list, "MESH")


def get_all_armature_obj(obj_list=None):
    return get_type_objs(obj_list, "ARMATURE")


def get_rest_bones(armature_obj):
    if armature_obj is None:
        return None, None
    else:
        rest_bones = []
        bones_idx_dict: "dict[str,int]" = {}
        for i, bone in enumerate(armature_obj.data.bones):
            rest_bones.append(armature_obj.matrix_world @ bone.head_local)
            bones_idx_dict[bone.name] = i
        rest_bones = np.stack(rest_bones, axis=0)
        # print(bones_idx_dict)
        return rest_bones, bones_idx_dict


def rename_maxiamo_bone(armature_obj):
    """Replace name like 'mixamorig10:xxx' to 'mixamorig:xxx, so that character can be correctly animated.'"""
    import re

    pattern = re.compile(r"mixamorig[0-9]+:")
    for bone in armature_obj.data.bones:
        bone.name = re.sub(pattern, "mixamorig:", bone.name)
    return armature_obj


def mesh_quads2tris(obj_list=None):
    if not obj_list:
        obj_list = bpy.context.scene.objects
    for obj in obj_list:
        if obj.type == "MESH":
            # obj.select_set(True)
            bpy.context.view_layer.objects.active = obj
            bpy.ops.object.mode_set(mode="EDIT")
            bpy.ops.mesh.quads_convert_to_tris(quad_method="BEAUTY", ngon_method="BEAUTY")
            bpy.ops.object.mode_set(mode="OBJECT")


if __name__ == "__main__":
    input_dir = "."
    output_dir = os.path.join(input_dir, "character_refined")
    os.makedirs(output_dir, exist_ok=True)
    character_list = sorted(glob(os.path.join(input_dir, "character", "*.fbx")))
    animation_list = sorted(glob(os.path.join(input_dir, "animation", "*.fbx")))
    unposed_list_file = "character_unposed.txt"
    if os.path.isfile(unposed_list_file):
        with open(unposed_list_file, "r") as f:
            unposed_list = f.readlines()
        unposed_list = [x.strip() for x in unposed_list if x.strip() != ""]
    else:
        unposed_list = None
    logfile = "character_bones.txt"

    bones_name_list = []
    for file in tqdm(character_list, dynamic_ncols=True):
        try:
            with HiddenPrints():
                remove_all()
                obj_list = load_file(file)
                armature_obj = get_all_armature_obj(obj_list)
                assert len(armature_obj) == 1, "Armature number is not 1"
                armature_obj = armature_obj[0]
                if unposed_list is None or file in unposed_list:
                    rename_maxiamo_bone(armature_obj)
                rest_bones, bones_idx_dict = get_rest_bones(armature_obj)

                # mesh_quads2tris(get_all_mesh_obj(obj_list))
                bpy.ops.object.select_all(action="DESELECT")
                for obj in obj_list:
                    obj.select_set(True)
                bpy.ops.export_scene.fbx(
                    filepath=os.path.join(output_dir, os.path.basename(file)),
                    check_existing=False,
                    use_selection=True,
                    use_triangles=True,  #!
                    add_leaf_bones=False,
                    bake_anim=False,
                    embed_textures=False,
                )
            bones_name_list += list(bones_idx_dict.keys())
            with open(logfile, "a") as f:
                f.write(f"{file}: {len(bones_idx_dict)}\n")
        except Exception as e:
            # raise
            tqdm.write(f"{file}: {e}")
            with open(logfile, "a") as f:
                f.write(f"{file}: {e}\n")

    # Calculate the coverage of each bone among all characters
    print(sorted(set(bones_name_list)))
    with HiddenPrints():
        remove_all()
        rest_bones, bones_idx_dict = get_rest_bones(get_all_armature_obj(load_file(animation_list[0]))[0])
    with open(logfile, "a") as f:
        for v in bones_idx_dict.keys():
            f.write(f"{v}: {bones_name_list.count(v)}/{len(character_list)}\n")