MotionCapture/Assets/Behaviours/BallStickBehaviour.cs

using System.Collections.Generic;
using Models;
using UnityEngine;
using Utils.PoseTransformHandlers;

namespace Behaviours
{
    public class BallStickBehaviour : MonoBehaviour
    {
        private readonly GameObject[] _nodes = new GameObject[PoseLandmarkType.MaxValue];
        private readonly List<Bond> _bonds = new List<Bond>();
        private readonly PoseTransform[] _poseTransforms = new PoseTransform[PoseLandmarkType.MaxValue];
        private readonly UdpListener _listener = new UdpListener();
        private const float Scale = 0.2f;

        // Start is called before the first frame update
        private void Start()
        {
            PoseTransform.AverageLength = 3;

            CreateNodes();
            CreateBonds();

            _listener.AddHandler(OnReceive);
            _listener.Connect(5000);
        }

        // Update is called once per frame
        private void Update()
        {
            var mid = _poseTransforms[PoseLandmarkType.LeftHip] + _poseTransforms[PoseLandmarkType.RightHip];
            mid = mid / 2;
            
            for (var i = 0; i < PoseLandmarkType.MaxValue; i++)
            {
                _nodes[i].transform.position = (_poseTransforms[i].ResultPosition - mid) * -5f;
            }

            foreach (var bond in _bonds)
            {
                bond.UpdateBond();
            }
        }

        private void OnDisable()
        {
            _listener.DisConnect();
        }

        private void OnReceive(List<PoseLandmark> landmarks)
        {
            foreach (var landmark in landmarks)
            {
                PoseTransform.UpdatePosition(ref _poseTransforms[landmark.Type.Value], landmark);
            }
        }

        private void CreateNodes()
        {
            for (var i = 0; i < PoseLandmarkType.MaxValue; i++)
            {
                var ball = GameObject.CreatePrimitive(PrimitiveType.Sphere);

                ball.transform.localScale = new Vector3(Scale, Scale, Scale);

                if (i <= 10)
                {
                    ball.GetComponent<Renderer>().material.color = Color.red; //给头部添加颜色
                }
                else if (i <= 22)
                {
                    ball.GetComponent<Renderer>().material.color = Color.blue; //给手部添加颜色
                }
                else if (i <= 32)
                {
                    ball.GetComponent<Renderer>().material.color = Color.green; //给脚部添加颜色
                }

                _nodes[i] = ball;
                _poseTransforms[i] = new PoseTransform(i, new AverageHandler());
            }
        }

        private void CreateBonds()
        {
            var headBonds = new[]
            {
                PoseLandmarkType.RightEar,
                PoseLandmarkType.RightEyeOuter,
                PoseLandmarkType.RightEye,
                PoseLandmarkType.RightEyeInner,
                PoseLandmarkType.Nose,
                PoseLandmarkType.LeftEyeInner,
                PoseLandmarkType.LeftEye,
                PoseLandmarkType.LeftEyeOuter,
                PoseLandmarkType.LeftEar
            };
            _bonds.AddRange(GenerateBondsList(headBonds));

            var monthBonds = new[]
            {
                PoseLandmarkType.MouthLeft,
                PoseLandmarkType.MouthRight,
            };
            _bonds.AddRange(GenerateBondsList(monthBonds));

            var leftArmBonds = new[]
            {
                PoseLandmarkType.LeftShoulder,
                PoseLandmarkType.LeftElbow,
                PoseLandmarkType.LeftWrist,
                PoseLandmarkType.LeftPinky,
                PoseLandmarkType.LeftIndex,
                PoseLandmarkType.LeftWrist,
                PoseLandmarkType.LeftThumb
            };
            _bonds.AddRange(GenerateBondsList(leftArmBonds));

            var rightArmBonds = new[]
            {
                PoseLandmarkType.RightShoulder,
                PoseLandmarkType.RightElbow,
                PoseLandmarkType.RightWrist,
                PoseLandmarkType.RightPinky,
                PoseLandmarkType.RightIndex,
                PoseLandmarkType.RightWrist,
                PoseLandmarkType.RightThumb
            };
            _bonds.AddRange(GenerateBondsList(rightArmBonds));

            var leftLegBonds = new[]
            {
                PoseLandmarkType.LeftShoulder,
                PoseLandmarkType.LeftHip,
                PoseLandmarkType.LeftKnee,
                PoseLandmarkType.LeftAnkle,
                PoseLandmarkType.LeftHeel,
                PoseLandmarkType.LeftFootIndex,
                PoseLandmarkType.LeftAnkle
            };
            _bonds.AddRange(GenerateBondsList(leftLegBonds));

            var rightLegBonds = new[]
            {
                PoseLandmarkType.RightShoulder,
                PoseLandmarkType.RightHip,
                PoseLandmarkType.RightKnee,
                PoseLandmarkType.RightAnkle,
                PoseLandmarkType.RightHeel,
                PoseLandmarkType.RightFootIndex,
                PoseLandmarkType.RightAnkle
            };
            _bonds.AddRange(GenerateBondsList(rightLegBonds));

            // 最后手动添加身体上的两条横线
            _bonds.Add(new Bond(
                _nodes[PoseLandmarkType.LeftShoulder],
                _nodes[PoseLandmarkType.RightShoulder],
                Scale));
            _bonds.Add(new Bond(
                _nodes[PoseLandmarkType.LeftHip],
                _nodes[PoseLandmarkType.RightHip],
                Scale));
        }

        /// <summary>
        /// 创建棍子列表
        /// </summary>
        /// <param name="nodes">需要连接起来的关键点 需要按顺序设置</param>
        /// <returns></returns>
        private List<Bond> GenerateBondsList(int[] nodes)
        {
            var bonds = new List<Bond>();

            for (var i = 0; i < nodes.Length - 1; i++)
            {
                bonds.Add(new Bond(
                    _nodes[nodes[i]],
                    _nodes[nodes[i + 1]],
                    Scale));
            }

            return bonds;
        }
    }
}