Main script from an extension to fix USD import of GLB characters in NVIDIA Omniverse (it is not fully working)

extension.py

# Made using: https://youtu.be/eGxV_PGNpOg
# Lessons learned
# - work out your package name first (it affects directory structure)
# - DeletePrims references Sdf which is not imported for you
import omni.ext
import omni.ui as ui
import omni.kit.commands
from pxr import Usd, Sdf, Gf


# Functions and vars are available to other extension as usual in python: `example.python_ext.some_public_function(x)`
def some_public_function(x: int):
    print("[ordinary] some_public_function was called with x: ", x)
    return x ** x


# Any class derived from `omni.ext.IExt` in top level module (defined in `python.modules` of `extension.toml`) will be
# instantiated when extension gets enabled and `on_startup(ext_id)` will be called. Later when extension gets disabled
# on_shutdown() is called.
class OrdinaryExtension(omni.ext.IExt):
    # ext_id is current extension id. It can be used with extension manager to query additional information, like where
    # this extension is located on filesystem.
    def on_startup(self, ext_id):
        self._window = ui.Window("VRM Import Cleanup v1", width=300, height=300)
        with self._window.frame:
            with ui.VStack():
                label = ui.Label("")

                def on_click():
                    self.clean_up_prim()
                    label.text = "clicked"

                def on_dump():
                    label.text = "dumped"
                    # TODO: Debugging: Print selected parts of hierarchy of stage to console
                    self.dump_stage()

                label.text = "dump"

                with ui.HStack():
                    ui.Button("Clean", clicked_fn=on_click)
                    ui.Button("Dump", clicked_fn=on_dump)

    def on_shutdown(self):
        print("[ordinary] ordinary shutdown")

    def clean_up_prim(self):

        # TODO: Could wrap in a big 'undo' wrapper so there are all undone together

        ctx = omni.usd.get_context()
        stage = ctx.get_stage()

        # Convert to SkelRoot type.
        # TODO: This is not a command, so Undo won't work on this one
        root_prim = stage.GetPrimAtPath('/World/Root')
        root_prim.SetTypeName('SkelRoot')

        # Move Skeleton directly under Root layer
        self.move_if_necessary(stage, '/World/Root/J_Bip_C_Hips0/Skeleton', '/World/Root/Skeleton')

        # Move meshes directly under Root layer, next to Skeleton
        self.move_if_necessary(stage, '/World/Root/J_Bip_C_Hips0/Face_baked', '/World/Root/Face_baked')
        self.move_if_necessary(stage, '/World/Root/J_Bip_C_Hips0/Body_baked', '/World/Root/Body_baked')
        self.move_if_necessary(stage, '/World/Root/J_Bip_C_Hips0/Hair001_baked', '/World/Root/Hair001_baked')

        # TODO: Delete - turns out not necessary - something recomputes it automatically.
        # self.add_parent_to_mesh_joint_list(stage, '/World/Root/Face_baked', 'Root')
        # self.add_parent_to_mesh_joint_list(stage, '/World/Root/Body_baked', 'Root')
        # self.add_parent_to_mesh_joint_list(stage, '/World/Root/Hair001_baked', 'Root')

        # Delete the dangling node (was old SkeRoot)
        self.delete_if_no_children(stage, '/World/Root/J_Bip_C_Hips0')

        # Patch the skeleton structure, if not done already. Add "Root" to the joint list.
        skeleton_prim = stage.GetPrimAtPath('/World/Root/Skeleton')
        if skeleton_prim:
            self.add_parent_to_skeleton_joint_list(skeleton_prim, 'Root')

    def move_if_necessary(self, stage, source_path, target_path):
        """
        If a prim exists at the source path, move it to the target path.
        Returns true if moved, false otherwise.
        """
        if stage.GetPrimAtPath(source_path):
            omni.kit.commands.execute(
                'MovePrim',
                path_from=source_path,
                path_to=target_path,
                keep_world_transform=False,
                destructive=False)
            return True
        return False

    def delete_if_no_children(self, stage, path):
        """
        Delete the prim at the specified path if it exists and has no children.
        Returns true if deleted, false otherwise.
        """
        print(path)
        prim = stage.GetPrimAtPath(path)
        if prim:
            print("found")
            if not prim.GetChildren():
                omni.kit.commands.execute(
                    'DeletePrims',
                    paths=[Sdf.Path(path)],
                    destructive=False)
                return True
        return False

    def add_parent_to_skeleton_joint_list(self, skeleton_prim: Usd.Prim, parent_name):
        """
        A skeleton has 3 attributes:
        - uniform matrix4d[] bindTransforms = [( (1, 0, -0, 0), ...]
        - uniform token[] joints = ["J_Bip_C_Hips", ...]
        - uniform matrix4d[] restTransforms = [( (1, 0, -0, 0), ...]

        In Omniverse Code etc, you can hover over the names in the "Raw USD Property" panel to get
        more documentation on the above properties.

        We need to insert the new parent at the front of the three lists, and prepend the name to the join paths.
        """

        # Get the attributes
        joints: Usd.Attribute = skeleton_prim.GetAttribute('joints')
        bindTransforms: Usd.Attribute = skeleton_prim.GetAttribute('bindTransforms')
        restTransforms: Usd.Attribute = skeleton_prim.GetAttribute('restTransforms')

        # If first join is the parent name already, nothing to do.
        if joints.Get()[0] == parent_name:
            return False

        # TODO: I use raw USD functions here, but there is also omni.kit.commands.execute("ChangeProperty",...)
        # if want undo...
        # https://docs.omniverse.nvidia.com/prod_kit/prod_kit/programmer_ref/usd/properties/set-attribute.html#omniverse-kit-commands
        joints.Set([parent_name] + [parent_name + '/' + jp for jp in joints.Get()])

        # Not 100% sure of this, but kinda works.
        # ((1, 0, -0, 0), (0, 1, 0, -0), (0, -0, 1, 0), (0, 0, 0, 1))
        unity_matrix = Gf.Matrix4d(1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1)
        if bindTransforms.IsValid():
            bindTransforms.Set([unity_matrix] + [x for x in bindTransforms.Get()])
        if restTransforms.IsValid():
            restTransforms.Set([unity_matrix] + [x for x in restTransforms.Get()])

    def add_parent_to_mesh_joint_list(self, stage, path, parent_name):
        """
        TODO: DELETE? THIS IS NOT ACTAULLY USED AT THE MOMENT.

        The meshes have paths to bones as well - add "Root" to their paths as well.
        """
        mesh_prim = stage.GetPrimAtPath(path)
        joints: Usd.Attribute = mesh_prim.GetAttribute('skel:joints')

        # Don't add 'Root' if its already been inserted. First value is empty by default.
        if joints.Get()[0] == "":
            joints.Set([parent_name if jp == "" else parent_name + '/' + jp for jp in joints.Get()])
            return True
        return False

    def dump_stage(self):
        """
        Traverse the tree of prims, printing out selected attribute information.
        Useful for debugging.
        """
        ctx = omni.usd.get_context()
        stage = ctx.get_stage()
        for prim in stage.Traverse():
            for attr in prim.GetAttributes():
                try:
                    if len(attr.Get()) >= 50:
                        print(attr.GetPath(), len(attr.Get()))
                except Exception:
                    pass