instance-rig 1.2.0__py3-none-any.whl

instancerig/_gltf.py

@@ -0,0 +1,169 @@
+from typing import List, Dict, Optional
+
+import numpy as np
+import open3d as o3d
+import pygltflib
+from volumesh import Volumesh
+from volumesh.GLTFMeshSequence import GLTFMeshSequence
+
+from instancerig import _math
+from instancerig.model.Joint import Joint
+from instancerig.model.Rig import Rig
+
+MESH_NAME = "mesh"
+
+
+def _append_joints(gltf: pygltflib.GLTF2, joint: Joint, joint_nodes: Dict[Joint, pygltflib.Node]) -> int:
+    child_ids: List[int] = []
+
+    relative_position = joint.relative_position
+    relative_rotation = joint.relative_rotation
+
+    for child in joint.children:
+        child_id = _append_joints(gltf, child, joint_nodes)
+        child_ids.append(child_id)
+
+    # Convert Euler angles to quaternion for GLTF
+    rx, ry, rz = np.radians(relative_rotation)
+    rotation_quaternion = _math.get_quaternion_from_euler(rx, ry, rz)
+
+    node = joint_nodes[joint]
+    node.children = child_ids
+    node.translation = list(relative_position)
+    node.rotation = rotation_quaternion
+    node.scale = [1.0, 1.0, 1.0]
+
+    node_id = gltf.nodes.index(node)
+    return node_id
+
+
+def export(output_path: str, mesh: o3d.geometry.TriangleMesh, rig: Rig, as_glb: bool = False):
+    gltf = export_to_gtlf(mesh, rig)
+
+    if as_glb:
+        gltf.save_binary(output_path)
+    else:
+        gltf.save_json(output_path)
+
+
+def export_to_gtlf(mesh: o3d.geometry.TriangleMesh, rig: Rig) -> pygltflib.GLTF2:
+    sequence = GLTFMeshSequence()
+    gltf = sequence.gltf
+
+    # append basic mesh infos
+    _append_mesh(sequence, mesh)
+    sequence_node = gltf.nodes[0]
+    mesh_node = gltf.nodes[1]
+    attributes = gltf.meshes[0].primitives[0].attributes
+
+    # make mesh root node
+    gltf.nodes.remove(sequence_node)
+
+    # add joint nodes to keep node-id the same as in skin weights (with delta)
+    joint_nodes: Dict[Joint, pygltflib.Node] = {}
+    joint_ids = []
+    for joint in list(rig.root_joint):
+        node = pygltflib.Node(name=joint.name)
+
+        node_id = len(gltf.nodes)
+        joint_ids.append(node_id)
+        gltf.nodes.append(node)
+
+        joint_nodes[joint] = node
+
+    # update joint nodes
+    root_joint_id = _append_joints(gltf, rig.root_joint, joint_nodes)
+    mesh_node.children = [root_joint_id]
+
+    # https://www.khronos.org/files/gltf20-reference-guide.pdf
+    # https://github.com/KhronosGroup/glTF-Tutorials/blob/master/gltfTutorial/gltfTutorial_019_SimpleSkin.md
+    # https://github.com/KhronosGroup/glTF-Tutorials/blob/master/gltfTutorial/gltfTutorial_020_Skins.md
+
+    # add bind matrices
+    gltf_bind_matrices = np.stack([j.inverse_bind_matrix.flatten(order="F") for j in list(rig.root_joint)])
+
+    gltf_bind_matrices = np.float32(gltf_bind_matrices)
+    inverse_bind_matrices_id = len(gltf.accessors)
+    _add_vector_data(sequence, gltf_bind_matrices, component_type=pygltflib.FLOAT, type=pygltflib.MAT4, target=None)
+
+    # add skin
+    skin = pygltflib.Skin(name="InstanceRigSkin", joints=joint_ids, inverseBindMatrices=inverse_bind_matrices_id)
+    skin_id = len(gltf.skins)
+    gltf.skins.append(skin)
+    mesh_node.skin = skin_id
+
+    # add joint_ids and weights
+    gltf_joint_ids = np.uint8(rig.joint_ids)
+    attributes.JOINTS_0 = len(gltf.accessors)
+    sequence._add_vector_data(gltf_joint_ids, component_type=pygltflib.UNSIGNED_BYTE, type=pygltflib.VEC4)
+
+    gltf_joint_weights = np.float32(rig.weights)
+    attributes.WEIGHTS_0 = len(gltf.accessors)
+    sequence._add_vector_data(gltf_joint_weights, component_type=pygltflib.FLOAT, type=pygltflib.VEC4)
+
+    return sequence.pack()
+
+
+def _append_mesh(sequence: GLTFMeshSequence, mesh: o3d.geometry.TriangleMesh):
+    # directly copied from volumesh -> todo: create own method in volumesh
+    name = MESH_NAME
+
+    # default arguments
+    points = np.float32(np.asarray(mesh.vertices))
+    triangles = np.uint32(np.asarray(mesh.triangles))
+
+    # optional arguments
+    colors = np.float32(np.asarray(mesh.vertex_colors))
+
+    # necessary to be GLTF compliant
+    mesh = mesh.normalize_normals()
+    normals = np.float32(np.asarray(mesh.vertex_normals))
+
+    # convert triangle_uvs into vertex_uvs
+    triangle_uvs = np.float32(np.asarray(mesh.triangle_uvs))
+    vertex_uvs = Volumesh._calculate_vertex_uvs(triangles, triangle_uvs) if len(triangle_uvs) > 0 else None
+
+    textures = [np.asarray(tex) for tex in mesh.textures if not tex.is_empty()]
+    texture = textures[0] if len(textures) > 0 else None
+
+    sequence.append_mesh(
+        points,
+        triangles,
+        colors,
+        normals,
+        vertex_uvs,
+        texture,
+        name=name,
+        compressed=False,
+        jpeg_textures=False,
+        jpeg_quality=100,
+    )
+
+
+def _add_vector_data(
+    sequence: GLTFMeshSequence,
+    array: np.ndarray,
+    component_type: int = pygltflib.FLOAT,
+    type: int = pygltflib.VEC3,
+    target: Optional[int] = pygltflib.ARRAY_BUFFER,
+):
+    array_blob = array.tobytes()
+    sequence.gltf.accessors.append(
+        pygltflib.Accessor(
+            bufferView=len(sequence.gltf.bufferViews),
+            componentType=component_type,
+            count=len(array),
+            type=str(type),
+            max=array.max(axis=0).tolist(),
+            min=array.min(axis=0).tolist(),
+        )
+    )
+    sequence.gltf.bufferViews.append(
+        pygltflib.BufferView(
+            buffer=0,
+            byteOffset=len(sequence.data),
+            byteLength=len(array_blob),
+            target=target,
+        )
+    )
+    sequence.data += array_blob

instancerig/_io.py

@@ -0,0 +1,12 @@
+import pickle
+from typing import Any
+
+
+def save_object(file_path: str, obj: Any):
+    with open(file_path, "wb") as file:
+        pickle.dump(obj, file)
+
+
+def load_object(file_path: str) -> Any:
+    with open(file_path, "rb") as file:
+        return pickle.load(file)

instancerig/_math.py

@@ -0,0 +1,163 @@
+import math
+import numpy as np
+
+from instancerig.model import Vector
+
+
+def interpolate(a: Vector, b: Vector, t: float) -> Vector:
+    return np.asarray([(t * b[0]) + ((1 - t) * a[0]), (t * b[1]) + ((1 - t) * a[1]), (t * b[2]) + ((1 - t) * a[2])])
+
+
+def distance(a: Vector, b: Vector) -> Vector:
+    return np.linalg.norm(np.asarray(a) - np.asarray(b), axis=0)
+
+
+def dot_product_angle(v1: Vector, v2: Vector) -> float:
+    if np.linalg.norm(v1) == 0 or np.linalg.norm(v2) == 0:
+        print("Zero magnitude vector!")
+    else:
+        vector_dot_product = np.dot(v1, v2)
+        arccos = np.arccos(vector_dot_product / (np.linalg.norm(v1) * np.linalg.norm(v2)))
+        angle = np.degrees(arccos)
+        return float(angle)
+    return 0
+
+
+def angle_between_points(a: Vector, b: Vector, axis: Vector, axis_id: int) -> float:
+    v = b - a
+    angle = dot_product_angle(axis, v)
+    return angle
+
+
+def normalize(data: np.ndarray) -> np.ndarray:
+    low, high = min(data), max(data)
+    if high - low == 0:
+        return np.zeros_like(data)
+    a = -low
+    b = 1 / (a + high)
+    return (data + a) * b
+
+
+def get_quaternion_from_euler(roll, pitch, yaw):
+    """
+    Convert an Euler angle to a quaternion.
+
+    Input
+      :param roll: The roll (rotation around x-axis) angle in radians.
+      :param pitch: The pitch (rotation around y-axis) angle in radians.
+      :param yaw: The yaw (rotation around z-axis) angle in radians.
+
+    Output
+      :return qx, qy, qz, qw: The orientation in quaternion [x,y,z,w] format
+    """
+    qx = np.sin(roll / 2) * np.cos(pitch / 2) * np.cos(yaw / 2) - np.cos(roll / 2) * np.sin(pitch / 2) * np.sin(yaw / 2)
+    qy = np.cos(roll / 2) * np.sin(pitch / 2) * np.cos(yaw / 2) + np.sin(roll / 2) * np.cos(pitch / 2) * np.sin(yaw / 2)
+    qz = np.cos(roll / 2) * np.cos(pitch / 2) * np.sin(yaw / 2) - np.sin(roll / 2) * np.sin(pitch / 2) * np.cos(yaw / 2)
+    qw = np.cos(roll / 2) * np.cos(pitch / 2) * np.cos(yaw / 2) + np.sin(roll / 2) * np.sin(pitch / 2) * np.sin(yaw / 2)
+
+    return [qx, qy, qz, qw]
+
+
+def get_rotation_matrix_from_euler(roll, pitch, yaw):
+    """
+    Convert Euler angles to a rotation matrix.
+    Assumes XYZ order.
+    """
+    R_x = np.array([[1, 0, 0], [0, math.cos(roll), -math.sin(roll)], [0, math.sin(roll), math.cos(roll)]])
+
+    R_y = np.array([[math.cos(pitch), 0, math.sin(pitch)], [0, 1, 0], [-math.sin(pitch), 0, math.cos(pitch)]])
+
+    R_z = np.array([[math.cos(yaw), -math.sin(yaw), 0], [math.sin(yaw), math.cos(yaw), 0], [0, 0, 1]])
+
+    R = np.dot(R_z, np.dot(R_y, R_x))
+    return R
+
+
+def get_euler_from_matrix(R):
+    """
+    Convert a rotation matrix to Euler angles (roll, pitch, yaw) in radians.
+    Assumes XYZ order.
+    """
+    sy = math.sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])
+    singular = sy < 1e-6
+
+    if not singular:
+        x = math.atan2(R[2, 1], R[2, 2])
+        y = math.atan2(-R[2, 0], sy)
+        z = math.atan2(R[1, 0], R[0, 0])
+    else:
+        x = math.atan2(-R[1, 2], R[1, 1])
+        y = math.atan2(-R[2, 0], sy)
+        z = 0
+
+    return np.array([x, y, z])
+
+
+def get_quaternion_from_matrix(R):
+    """
+    Convert a rotation matrix to a quaternion [x, y, z, w].
+    """
+    trace = np.trace(R)
+    if trace > 0:
+        S = np.sqrt(trace + 1.0) * 2
+        qw = 0.25 * S
+        qx = (R[2, 1] - R[1, 2]) / S
+        qy = (R[0, 2] - R[2, 0]) / S
+        qz = (R[1, 0] - R[0, 1]) / S
+    elif (R[0, 0] > R[1, 1]) and (R[0, 0] > R[2, 2]):
+        S = np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2]) * 2
+        qw = (R[2, 1] - R[1, 2]) / S
+        qx = 0.25 * S
+        qy = (R[0, 1] + R[1, 0]) / S
+        qz = (R[0, 2] + R[2, 0]) / S
+    elif R[1, 1] > R[2, 2]:
+        S = np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2]) * 2
+        qw = (R[0, 2] - R[2, 0]) / S
+        qx = (R[0, 1] + R[1, 0]) / S
+        qy = 0.25 * S
+        qz = (R[1, 2] + R[2, 1]) / S
+    else:
+        S = np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1]) * 2
+        qw = (R[1, 0] - R[0, 1]) / S
+        qx = (R[0, 2] + R[2, 0]) / S
+        qy = (R[1, 2] + R[2, 1]) / S
+        qz = 0.25 * S
+
+    return [qx, qy, qz, qw]
+
+
+def get_alignment_matrix(target_vector, source_vector=np.array([0, 1, 0])):
+    """
+    Calculate rotation matrix to align source_vector to target_vector.
+    """
+    target_vector = target_vector / (np.linalg.norm(target_vector) + 1e-6)
+    source_vector = source_vector / (np.linalg.norm(source_vector) + 1e-6)
+
+    v = np.cross(source_vector, target_vector)
+    c = np.dot(source_vector, target_vector)
+    s = np.linalg.norm(v)
+
+    if s < 1e-6:
+        # Vectors are parallel
+        if c > 0:
+            return np.eye(3)
+        else:
+            # Vectors are opposite
+            # Find an arbitrary perpendicular vector to rotate 180 degrees
+            if np.abs(source_vector[0]) < 0.1:
+                perp = np.array([1, 0, 0])
+            else:
+                perp = np.array([0, 1, 0])
+            axis = np.cross(source_vector, perp)
+            axis = axis / np.linalg.norm(axis)
+            # Rotation of 180 degrees around axis
+            # Rodrigues formula for 180 deg: I + 2 * K^2
+            # K is skew symmetric matrix of axis
+            # Or just construct rotation matrix manually
+            # R = 2 * outer(axis, axis) - I
+            return 2 * np.outer(axis, axis) - np.eye(3)
+
+    kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
+
+    R = np.eye(3) + kmat + np.dot(kmat, kmat) * ((1 - c) / (s**2))
+    return R

instancerig/_pose_correction.py

@@ -0,0 +1,154 @@
+from typing import Dict, Tuple
+
+import numpy as np
+import open3d as o3d
+
+from instancerig import _math
+from instancerig.model.Joint import Joint
+from instancerig.model.Rig import Rig
+
+
+class TPoseConverter:
+    def __init__(self, rig: Rig, mesh: o3d.geometry.TriangleMesh):
+        self.rig = rig
+        self.mesh = mesh
+        self.joints = {j.name: j for j in list(rig.root_joint)}
+
+        # Define the coordinate system.
+        # Z as Forward, Y as Up, X as Right.
+        self.forward = np.array([0, 0, 1])
+        self.up = np.array([0, 1, 0])
+        self.right = np.array([1, 0, 0])
+
+    def convert_to_t_pose(self) -> Tuple[o3d.geometry.TriangleMesh, Rig]:
+        parameters = self._calculate_correction_parameters()
+
+        # Define rotations (in degrees) for joints to reach T-pose.
+        # Rotation around Z-axis.
+        rotations = {
+            "rightShoulder": np.array([0, 0, parameters["ShoulderAngleShiftRight"]]),
+            "leftShoulder": np.array([0, 0, -parameters["ShoulderAngleShiftLeft"]]),
+            "rightHip": np.array([0, 0, parameters["HipAngleShiftRight"]]),
+            "leftHip": np.array([0, 0, -parameters["HipAngleShiftLeft"]]),
+        }
+
+        self._apply_t_pose(rotations)
+
+        return self.mesh, self.rig
+
+    def _calculate_correction_parameters(self) -> Dict[str, float]:
+        def get_pos(name):
+            if name not in self.joints:
+                return np.array([0.0, 0.0, 0.0])
+            return self.joints[name].world_position
+
+        def project_2d(v):
+            # Project vector onto the plane perpendicular to forward (the XY plane)
+            return v - np.dot(v, self.forward) * self.forward
+
+        def angle_between(a_name, b_name, c_name):
+            a, b, c = get_pos(a_name), get_pos(b_name), get_pos(c_name)
+
+            # Project points to 2D (XY plane) to ignore flexion/extension (forward/backward bending)
+            # This ensures we only correct the abduction/adduction (side split)
+            a = project_2d(a)
+            b = project_2d(b)
+            c = project_2d(c)
+
+            v1 = a - b
+            v2 = c - b
+            return _math.dot_product_angle(v1, v2)
+
+        def angle_to_up(start, end):
+            a = project_2d(end) - project_2d(start)
+            up_proj = project_2d(self.up)
+            return _math.dot_product_angle(up_proj, a)
+
+        # Calculate angles relative to the desired T-pose
+        right_shoulder_angle = angle_between("rightHip", "rightShoulder", "rightElbow") - 90
+        left_shoulder_angle = angle_between("leftHip", "leftShoulder", "leftElbow") - 90
+
+        # Hip angle calculation:
+        # 180 - angle gives the positive correction needed to close the legs (Adduction).
+        # We use projected angles to avoid over-correction due to forward knee bend.
+        right_hip_angle = angle_between("rightShoulder", "rightHip", "rightKnee") - 180
+        left_hip_angle = angle_between("leftShoulder", "leftHip", "leftKnee") - 180
+
+        # Calculate torso tilt correction
+        right_side_angle = 180.0 - angle_to_up(get_pos("rightShoulder"), get_pos("rightHip"))
+        left_side_angle = 180.0 - angle_to_up(get_pos("leftShoulder"), get_pos("leftHip"))
+
+        return {
+            "ShoulderAngleShiftRight": right_shoulder_angle + right_side_angle,
+            "ShoulderAngleShiftLeft": left_shoulder_angle + left_side_angle,
+            "HipAngleShiftRight": right_hip_angle + right_side_angle,
+            "HipAngleShiftLeft": left_hip_angle + left_side_angle,
+        }
+
+    def _apply_t_pose(self, joint_rotations: Dict[str, np.ndarray]):
+        node_transforms = {}  # joint_name -> 4x4 matrix
+
+        def traverse(joint: Joint, parent_m: np.ndarray):
+            # Correction rotation for this joint
+            rel_euler = joint_rotations.get(joint.name, np.array([0.0, 0.0, 0.0]))
+            rx, ry, rz = np.radians(rel_euler)
+            R = _math.get_rotation_matrix_from_euler(rx, ry, rz)
+
+            # Transformation around joint position: T * R * T^-1
+            pos = joint.world_position
+            T = np.eye(4)
+            T[:3, 3] = pos
+            T_inv = np.eye(4)
+            T_inv[:3, 3] = -pos
+
+            R_4x4 = np.eye(4)
+            R_4x4[:3, :3] = R
+
+            local_m = T @ R_4x4 @ T_inv
+
+            # Cumulative transform
+            total_m = parent_m @ local_m
+            node_transforms[joint.name] = total_m
+
+            for child in joint.children:
+                traverse(child, total_m)
+
+        traverse(self.rig.root_joint, np.eye(4))
+
+        # Apply Linear Blend Skinning (LBS) to mesh vertices
+        vertices = np.asarray(self.mesh.vertices)
+        new_vertices = np.zeros_like(vertices)
+
+        joint_list = list(self.rig.root_joint)
+        joint_names = [j.name for j in joint_list]
+        matrices = np.array([node_transforms[name] for name in joint_names])
+
+        for i in range(len(vertices)):
+            v = np.append(vertices[i], 1.0)
+            v_new = np.zeros(4)
+            for j in range(4):  # 4 joints per vertex
+                joint_idx = self.rig.joint_ids[i, j]
+                weight = self.rig.weights[i, j]
+                if weight > 0:
+                    v_new += weight * (matrices[joint_idx] @ v)
+            new_vertices[i] = v_new[:3]
+
+        self.mesh.vertices = o3d.utility.Vector3dVector(new_vertices)
+
+        # Update joint positions and rotations in the rig
+        for joint in joint_list:
+            m = node_transforms[joint.name]
+
+            # Update position
+            pos_4 = np.append(joint.world_position, 1.0)
+            joint.world_position = (m @ pos_4)[:3]
+
+            # Update rotation
+            R_applied = m[:3, :3]
+            rx, ry, rz = np.radians(joint.world_rotation)
+            R_old = _math.get_rotation_matrix_from_euler(rx, ry, rz)
+            R_new = R_applied @ R_old
+            joint.world_rotation = np.degrees(_math.get_euler_from_matrix(R_new))
+
+        # Recompute normals
+        self.mesh.compute_vertex_normals()

instancerig/_preview.py

@@ -0,0 +1,129 @@
+from typing import List, Iterable, Tuple
+
+import numpy as np
+import open3d as o3d
+from muke.model.KeyPoint3 import KeyPoint3
+
+from instancerig._math import get_alignment_matrix
+from instancerig.model import Vector
+from instancerig.model.Joint import Joint
+
+COLOR_SEQUENCE = [
+    (230, 25, 75),
+    (60, 180, 75),
+    (255, 225, 25),
+    (0, 130, 200),
+    (245, 130, 48),
+    (145, 30, 180),
+    (70, 240, 240),
+    (240, 50, 230),
+    (210, 245, 60),
+    (250, 190, 212),
+    (0, 128, 128),
+    (220, 190, 255),
+    (170, 110, 40),
+    (255, 250, 200),
+    (128, 0, 0),
+    (170, 255, 195),
+    (128, 128, 0),
+    (255, 215, 180),
+    (255, 255, 255),
+]
+
+NICE_COLOR_SEQUENCE = [
+    [51, 87, 255],  # Blue
+    [87, 255, 51],  # Green
+    [255, 102, 51],  # Red-Orange
+    [255, 51, 255],  # Pink
+    [255, 255, 51],  # Yellow
+    [255, 51, 102],  # Coral
+    [51, 153, 255],  # Sky Blue
+    [153, 51, 255],  # Purple
+    [51, 255, 255],  # Cyan
+    [255, 102, 51],  # Orange
+    [51, 255, 153],  # Mint Green
+    [204, 51, 255],  # Lavender
+    [102, 255, 51],  # Lime Green
+    [255, 51, 204],  # Magenta
+    [51, 204, 255],  # Light Blue
+    [255, 204, 51],  # Gold
+    [51, 255, 204],  # Turquoise
+    [255, 153, 51],  # Tangerine
+    [51, 255, 51],  # Bright Green
+    [255, 51, 255],  # Orchid
+]
+
+PREVIEW_WIDTH = 720
+PREVIEW_HEIGHT = 1280
+
+
+def draw_cone(
+    bottom_center: Iterable[float],
+    top_position: Iterable[float],
+    color: Tuple[float, float, float] = (0.6, 0.6, 0.9),
+    radius: float = 0.007,
+    tip_offset: float = 5e-3,
+) -> o3d.geometry.TriangleMesh:
+    bottom: np.ndarray = np.asarray(bottom_center, dtype=float)
+    top: np.ndarray = np.asarray(top_position, dtype=float)
+
+    direction: np.ndarray = top - bottom
+    length: float = float(np.linalg.norm(direction))
+
+    cone: o3d.geometry.TriangleMesh = o3d.geometry.TriangleMesh.create_cone(radius=radius, height=length + 1e-6)
+
+    R: np.ndarray = get_alignment_matrix(target_vector=direction, source_vector=np.array([0.0, 0.0, 1.0], dtype=float))
+    cone.rotate(R, center=(0.0, 0.0, 0.0))
+
+    if length > 1e-12:
+        cone.translate(bottom + tip_offset * (direction / length))
+    else:
+        cone.translate(bottom)
+
+    cone.paint_uniform_color(np.asarray(color, dtype=float))
+    return cone
+
+
+def preview_3d_keypoints(mesh: o3d.geometry.TriangleMesh, keypoints: List[KeyPoint3], title: str = "3D KeyPoints"):
+    meshes = []
+
+    wireframe = o3d.geometry.LineSet.create_from_triangle_mesh(mesh)
+    meshes.append(wireframe)
+
+    for kp in keypoints:
+        marker: o3d.geometry.TriangleMesh = o3d.geometry.TriangleMesh.create_sphere(0.025)
+        marker.paint_uniform_color([1, 0.706, 0])
+        marker.translate(kp.predicted_position)
+        meshes.append(marker)
+
+    o3d.visualization.draw_geometries(meshes, window_name=title, width=PREVIEW_WIDTH, height=PREVIEW_HEIGHT)
+
+
+def render_joints(
+    mesh: o3d.geometry.TriangleMesh, root_joint: Joint, markers: List[Vector] = None, title: str = "Joints"
+):
+    meshes = []
+    if markers is None:
+        markers = []
+
+    wireframe = o3d.geometry.LineSet.create_from_triangle_mesh(mesh)
+    meshes.append(wireframe)
+
+    for joint in root_joint:
+        joint_marker: o3d.geometry.TriangleMesh = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1)
+        joint_marker.translate(joint.world_position)
+        R = joint_marker.get_rotation_matrix_from_xyz(np.radians(joint.world_rotation))
+        joint_marker.rotate(R, center=joint.world_position)
+        meshes.append(joint_marker)
+
+        # render bones
+        for child in joint.children:
+            meshes.append(draw_cone(joint.world_position, child.world_position, color=[1.0, 0.2, 0.2], radius=0.02))
+
+    for marker_position in markers:
+        marker: o3d.geometry.TriangleMesh = o3d.geometry.TriangleMesh.create_sphere(radius=0.025)
+        marker.paint_uniform_color([1, 0.2, 0])
+        marker.translate(marker_position)
+        meshes.append(marker)
+
+    o3d.visualization.draw_geometries(meshes, window_name=title, width=PREVIEW_WIDTH, height=PREVIEW_HEIGHT)

instancerig/_utils.py

@@ -0,0 +1,58 @@
+import os
+import platform
+import time
+from pathlib import Path
+
+
+def current_millis() -> int:
+    return time.time_ns() // 1_000_000
+
+
+def get_cache_dir() -> Path:
+    if cache_dir_env := os.getenv("INSTANCERIG_CACHE_DIR"):
+        path = Path(cache_dir_env)
+    else:
+        if platform.system() == "Windows":
+            path = Path.home() / "AppData" / "Local" / "instancerig"
+        else:
+            path = Path.home() / ".cache" / "instancerig"
+
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+class Watch:
+    def __init__(self, name: str = "Watch"):
+        self.name = name
+        self.start_time: int = 0
+        self.end_time: int = 0
+        self.running = False
+
+    def start(self):
+        self.start_time = current_millis()
+        self.running = True
+
+    def stop(self):
+        self.end_time = current_millis()
+        self.running = False
+
+    def elapsed(self) -> int:
+        if self.running:
+            return current_millis() - self.start_time
+        return self.end_time - self.start_time
+
+    def print(self):
+        print(f"{self.name}: {self.elapsed_str()}")
+
+    def elapsed_str(self, correction: int = 0) -> str:
+        delta = self.elapsed() - correction
+        time_info = time.strftime("%Hh %Mm %Ss {}ms".format(delta % 1000), time.gmtime(delta / 1000.0))
+        return time_info
+
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop()
+        self.print()