MotionController.cs

using System.Collections;
using System.Collections.Generic;
using System.IO;
using UnityEngine;

/// <summary>
/// Plays continuous SMPL human motion sequences loaded from JSON files in the Unity Resources folder.
/// Handles frame timing based on recorded framerate, root translation, root orientation, and joint rotations.
/// </summary>
public class MotionController : MonoBehaviour
{
    [Header("Playback & File Loading")]
    [Tooltip("Name of the motion data JSON file (without .json extension) to load from the Resources folder.")]
    public string motionDataFilename = "01";

    [Tooltip("Set to true in the Inspector during play mode to reload motion data and restart playback from frame 0.")]
    public bool loadNewMotionData = false;

    [Header("Target Model & Origin")]
    [Tooltip("Target SMPL model GameObject. If unassigned, defaults to the GameObject this component is attached to.")]
    public GameObject smplModelInstance;

    [Header("Root Transform Controls")]
    [Tooltip("Whether to apply root (pelvis) translation from the motion sequence.")]
    public bool rootTranslation = true;

    [Tooltip("Whether to apply root (pelvis) rotation from the motion sequence.")]
    public bool rootRotation = true;

    public enum BodyModelShape
    {
        Neutral,
        Male,
        Female
    }

    [Header("Body Shape Compensation")]
    [Tooltip("Body model shape configuration for the 3D mesh (Neutral, Male, or Female). Used for height scaling compensation.")]
    public BodyModelShape meshShape = BodyModelShape.Male;

    [Tooltip("Body model shape configuration assumed by the motion data (Neutral, Male, or Female). Used for height scaling compensation.")]
    public BodyModelShape motionShape = BodyModelShape.Neutral;

    private MotionData _motionData;
    private Vector3[] _translations;
    private Quaternion[,] _rotations;
    private int _currentFrame = 0;
    private float _timer = 0f;

    [System.Serializable]
    public class MotionData
    {
        public int n_frames;
        public int n_joints;
        public float framerate;
        public List<float> translations;  // [frame][0][x,y,z]
        public List<float> rotations;   // [frame][joint][x,y,z,w]
    }

    private static readonly Dictionary<BodyModelShape, float> ModelShapeHeight = new Dictionary<BodyModelShape, float>
    {
        { BodyModelShape.Neutral, 1.717f },
        { BodyModelShape.Male, 1.789f },
        { BodyModelShape.Female, 1.657f }
    };

    private static readonly Dictionary<int, string> JOINT_NAMES = new Dictionary<int, string>
    {
        { 0, "pelvis" },
        { 1, "left_hip" },
        { 2, "right_hip" },
        { 3, "spine1" },
        { 4, "left_knee" },
        { 5, "right_knee" },
        { 6, "spine2" },
        { 7, "left_ankle" },
        { 8, "right_ankle" },
        { 9, "spine3" },
        { 10, "left_foot" },
        { 11, "right_foot" },
        { 12, "neck" },
        { 13, "left_collar" },
        { 14, "right_collar" },
        { 15, "head" },
        { 16, "left_shoulder" },
        { 17, "right_shoulder" },
        { 18, "left_elbow" },
        { 19, "right_elbow" },
        { 20, "left_wrist" },
        { 21, "right_wrist" },
        { 22, "left_hand" },
        { 23, "right_hand" }
    };

    private Dictionary<string, Transform> _jointTransforms = new Dictionary<string, Transform>();

    private Vector3 _pelvisInitialPosition;
    private Quaternion _pelvisInitialRotation;
    private Transform _modelOrigin;

    void Start()
    {
        // If smplModelInstance is not set, look for the GameObject to which the script is attached
        if (smplModelInstance == null)
            smplModelInstance = gameObject;
        _modelOrigin = smplModelInstance.transform;

        // Initialize jointTransforms dictionary
        foreach (var joint in JOINT_NAMES.Values)
        {
            Transform jointTransform = FindChildRecursively(smplModelInstance.transform, joint);
            if (jointTransform != null)
            {
                _jointTransforms[joint] = jointTransform;
            }
            else
            {
                Debug.LogWarning($"Joint {joint} not found in the SMPL model.");
            }
        }

        // Save initial pelvis position and rotation
        if (_jointTransforms.TryGetValue("pelvis", out Transform pelvisTransform))
        {
            _pelvisInitialPosition = pelvisTransform.localPosition;
            _pelvisInitialRotation = pelvisTransform.localRotation;
        }

        // Load motion data
        LoadMotionData(motionDataFilename);
    }

    void LoadMotionData(string filename)
    {
        TextAsset jsonFile = Resources.Load<TextAsset>(filename);
        if (jsonFile != null)
        {
            _motionData = JsonUtility.FromJson<MotionData>(jsonFile.text);
            if (_motionData.framerate <= 0)
            {
                Debug.LogWarning($"Framerate in '{filename}' is <= 0 ({_motionData.framerate}). Defaulting to 25 FPS.");
                _motionData.framerate = 25f;
            }
            Debug.Log($"Loaded {_motionData.n_frames} frames of motion data");
            Debug.Log($"Framerate: {_motionData.framerate}");
            Debug.Log($"Number of joints: {_motionData.n_joints}");

            _translations = new Vector3[_motionData.n_frames];
            for (int i = 0; i < _motionData.n_frames; i++)
            {
                // Convert to Unity coordinate system
                _translations[i] = new Vector3(
                    -_motionData.translations[i * 3],
                    _motionData.translations[i * 3 + 2],
                    -_motionData.translations[i * 3 + 1]
                );
            }

            _rotations = new Quaternion[_motionData.n_frames, _motionData.n_joints];
            for (int i = 0; i < _motionData.n_frames; i++)
            {
                for (int j = 0; j < _motionData.n_joints; j++)
                {
                    int baseIndex = (i * _motionData.n_joints + j) * 4;
                    // Convert to Unity coordinate system
                    _rotations[i, j] = new Quaternion(
                        _motionData.rotations[baseIndex],
                        -_motionData.rotations[baseIndex + 1],
                        -_motionData.rotations[baseIndex + 2],
                        _motionData.rotations[baseIndex + 3]
                    );
                }
            }

            Debug.Log("Loaded motion data");
        }
        else
            Debug.LogError($"Could not load '{filename}' from Resources folder");
    }

    void Update()
    {
        if (loadNewMotionData)
        {
            _currentFrame = 0;
            _timer = 0f;
            LoadMotionData(motionDataFilename);
            loadNewMotionData = false;
            Debug.Log("New motion data loaded");
        }
        if (_motionData == null || _motionData.n_frames == 0)
            return;

        _timer += Time.deltaTime;
        float targetInterval = (_motionData.framerate > 0) ? 1f / _motionData.framerate : 1f / 25f;
        if (_timer >= targetInterval)
        {
            _timer -= targetInterval;
            UpdateMotion();
            _currentFrame = (_currentFrame + 1) % _motionData.n_frames;
        }
    }

    void UpdateMotion()
    {
        for (int i = 0; i < _motionData.n_joints; i++)
        {
            if (JOINT_NAMES.TryGetValue(i, out string jointName) && _jointTransforms.TryGetValue(jointName, out Transform jointTransform))
            {
                if (jointName == "pelvis")
                {
                    if (rootTranslation)
                    {
                        float yScale = ModelShapeHeight[meshShape] / ModelShapeHeight[motionShape];
                        Vector3 currentTrans = new Vector3(
                            _translations[_currentFrame].x,
                            _translations[_currentFrame].y * yScale,
                            _translations[_currentFrame].z
                        );
                        jointTransform.position = _modelOrigin.TransformPoint(currentTrans);
                    }
                    else
                        jointTransform.localPosition = _pelvisInitialPosition;
                    if (rootRotation)
                        jointTransform.localRotation = _pelvisInitialRotation * Quaternion.Euler(-90f, 0f, 0f) * _rotations[_currentFrame, i];
                    else
                        jointTransform.localRotation = _pelvisInitialRotation;
                }
                else
                    jointTransform.localRotation = _rotations[_currentFrame, i];
            }

        }
    }

    private static Transform FindChildRecursively(Transform parent, string name)
    {
        foreach (Transform child in parent)
        {
            if (child.name == name)
                return child;
            Transform result = FindChildRecursively(child, name);
            if (result != null)
                return result;
        }
        return null;
    }
}