RuntimeSkeletonVisualizer.cs

using System.Collections.Generic;
using UnityEngine;

/// <summary>
/// Generates and updates a 3D visual skeleton at runtime using spheres for joints and tapered cylinders
/// for bones. Automatically discovers bone hierarchy from a SkinnedMeshRenderer.
/// </summary>
public class RuntimeSkeletonVisualizer : MonoBehaviour
{
    [Header("Bone Settings")]
    [Tooltip("The root bone of the skeleton hierarchy. If autoDetectBones is true, this is resolved automatically from the SkinnedMeshRenderer.")]
    public Transform rootBone;
    [Tooltip("Whether to skip the root armature bone and start visualization at the first child joint (pelvis).")]
    public bool skipRootBone = true;

    [Header("Dimensions")]
    [Tooltip("Radius of the joint spheres in millimeters.")]
    [Range(0f, 100f)]
    public float jointRadius = 45f;
    [Tooltip("Width/diameter of the bone cylinders in millimeters.")]
    [Range(0f, 80f)]
    public float boneWidth = 30f;
    [Tooltip("Cone tapering factor between parent and child joint (0 = full point, 1 = uniform cylinder).")]
    [Range(0.0f, 1.0f)]
    public float coneFactor = 0.25f;

    [Header("Materials & Colors")]
    [Tooltip("Material applied to joint spheres. Defaults to standard blue if unassigned.")]
    public Material jointMaterial;
    [Tooltip("Material applied to bone cylinders. Defaults to standard white if unassigned.")]
    public Material boneMaterial;
    [Tooltip("Fallback color for joints if jointMaterial is not assigned.")]
    public Color jointColor = Color.blue;
    [Tooltip("Fallback color for bones if boneMaterial is not assigned.")]
    public Color boneColor = Color.white;

    [Header("Detection")]
    [Tooltip("Automatically detects bone transforms from the SkinnedMeshRenderer attached to this model or its children.")]
    public bool autoDetectBones = true;

    private List<GameObject> visualObjects = new List<GameObject>();
    private Transform[] bones;

    private GameObject rigParent;

    void Start()
    {
        rigParent = new GameObject("RigParent");
        rigParent.transform.SetParent(transform);

        if (autoDetectBones)
        {
            SkinnedMeshRenderer smr = GetComponentInChildren<SkinnedMeshRenderer>();
            if (smr == null)
            {
                Debug.LogError("No SkinnedMeshRenderer found in children.");
                return;
            }
            else
            {
                bones = smr.bones;
                if (bones.Length > 0 && rootBone == null)
                {
                    // Find root bone
                    HashSet<Transform> boneSet = new HashSet<Transform>(bones);
                    foreach (Transform bone in bones)
                    {
                        if (bone.parent == null || !boneSet.Contains(bone.parent))
                        {
                            rootBone = bone;
                            break;
                        }
                    }
                    if (rootBone != null && skipRootBone && rootBone.childCount > 0)
                    {
                        if (rootBone.childCount > 1)
                        {
                            Debug.LogWarning("Root bone: " + rootBone.name + " has " + rootBone.childCount + " children.");
                            foreach (Transform child in rootBone)
                            {
                                Debug.LogWarning("Child name: " + child.name);
                            }
                            Debug.LogWarning("Using the first child as root bone.");
                        }
                        rootBone = rootBone.GetChild(0);
                    }
                }
            }
        }

        if (rootBone == null)
        {
            Debug.LogError("Root bone not assigned.");
            return;
        }

        // Create default materials if not assigned
        Shader defaultShader = Shader.Find("Standard") ?? Shader.Find("Universal Render Pipeline/Lit") ?? Shader.Find("Mobile/Diffuse");
        if (jointMaterial == null)
        {
            Debug.LogWarning("Joint material not assigned. Creating default material.");
            jointMaterial = new Material(defaultShader);
            jointMaterial.color = jointColor;
        }
        if (boneMaterial == null)
        {
            Debug.LogWarning("Bone material not assigned. Creating default material.");
            boneMaterial = new Material(defaultShader);
            boneMaterial.color = boneColor;
        }

        CreateSkeletonVisualization(rootBone);
    }

    void Update()
    {
        UpdateSkeletonVisualization();
    }

    void CreateSkeletonVisualization(Transform bone, Transform parentBone = null)
    {
        // Create a sphere to represent the joint
        GameObject jointObj = GameObject.CreatePrimitive(PrimitiveType.Sphere);
        Collider jointCol = jointObj.GetComponent<Collider>();
        if (jointCol != null) Destroy(jointCol);
        jointObj.GetComponent<Renderer>().shadowCastingMode = UnityEngine.Rendering.ShadowCastingMode.Off;
        jointObj.name = "Joint_" + bone.name;
        jointObj.transform.SetParent(rigParent.transform);
        jointObj.transform.position = bone.position;
        jointObj.transform.localScale = Vector3.one * jointRadius / 1000;
        jointObj.GetComponent<Renderer>().sharedMaterial = jointMaterial;
        visualObjects.Add(jointObj);

        // Create a cone to represent the bone if there is a parent
        if (parentBone != null)
        {
            GameObject boneObj = CreateBoneVisualization(bone, parentBone);
            visualObjects.Add(boneObj);
        }

        // Create visualizations for all children
        for (int i = 0; i < bone.childCount; i++)
        {
            Transform childBone = bone.GetChild(i);
            if (IsValidBone(bone, childBone))
            {
                CreateSkeletonVisualization(childBone, bone);
            }
        }
    }

    private bool IsValidBone(Transform bone, Transform childBone)
    {
        bool isValidBone = true;
        if (bones != null && bones.Length > 0)
        {
            isValidBone = System.Array.Exists(bones, b => b == childBone);
        }

        if (bone.name.Contains("wrist") || bone.name.Contains("head") || bone.name.Contains("Carpus"))
        {
            isValidBone = false;
        }

        return isValidBone;
    }

    GameObject CreateBoneVisualization(Transform bone, Transform parentBone)
    {
        GameObject boneObj = GameObject.CreatePrimitive(PrimitiveType.Cylinder);
        Collider boneCol = boneObj.GetComponent<Collider>();
        if (boneCol != null) Destroy(boneCol);
        boneObj.GetComponent<Renderer>().shadowCastingMode = UnityEngine.Rendering.ShadowCastingMode.Off;
        boneObj.name = "Bone_" + parentBone.name + "_to_" + bone.name;
        boneObj.transform.SetParent(rigParent.transform);

        // Set bone position and rotation
        Vector3 direction = bone.position - parentBone.position;
        float distance = direction.magnitude;

        boneObj.transform.position = parentBone.position + direction * 0.5f;
        boneObj.transform.up = direction.normalized;
        boneObj.transform.localScale = new Vector3(boneWidth / 1000, distance * 0.5f, boneWidth / 1000);

        // Adjust the cylinder to look like a cone
        MeshFilter meshFilter = boneObj.GetComponent<MeshFilter>();
        Mesh mesh = meshFilter.mesh;
        Vector3[] vertices = mesh.vertices;
        for (int i = 0; i < vertices.Length; i++)
        {
            if (vertices[i].y > 0)
            {
                vertices[i].x *= coneFactor;
                vertices[i].z *= coneFactor;
            }
        }
        mesh.vertices = vertices;
        mesh.RecalculateNormals();

        boneObj.GetComponent<Renderer>().sharedMaterial = boneMaterial;

        return boneObj;
    }

    void UpdateSkeletonVisualization()
    {
        int visualIndex = 0;
        UpdateBoneVisualization(rootBone, null, ref visualIndex);
    }

    void UpdateBoneVisualization(Transform bone, Transform parentBone, ref int visualIndex)
    {
        // Update joint position
        GameObject jointObj = visualObjects[visualIndex++];
        jointObj.transform.position = bone.position;
        jointObj.transform.localScale = Vector3.one * jointRadius / 1000;

        // Update bone position and rotation if there is a parent
        if (parentBone != null)
        {
            GameObject boneObj = visualObjects[visualIndex++];

            Vector3 direction = bone.position - parentBone.position;
            float distance = direction.magnitude;

            boneObj.transform.position = parentBone.position + direction * 0.5f;
            boneObj.transform.up = direction.normalized;
            boneObj.transform.localScale = new Vector3(boneWidth / 1000, distance * 0.5f, boneWidth / 1000);
        }

        // Update children recursively
        for (int i = 0; i < bone.childCount; i++)
        {
            Transform childBone = bone.GetChild(i);
            if (IsValidBone(bone, childBone))
            {
                UpdateBoneVisualization(childBone, bone, ref visualIndex);
            }
        }
    }

    void OnDestroy()
    {
        if (rigParent != null)
        {
            Destroy(rigParent);
        }
        visualObjects.Clear();
    }
}