Parser for 3D Keypoints

2024-09-17 09:35 UTC gpt-4o Open in ChatGPT ↗

To create a parser for .pkl files containing 3D keypoints in the COCO format, obtained from the AIST++ Dance Video Database, you can utilize Python’s pickle module to read the .pkl files and extract the 3D keypoints. This parser will be designed to work with Unity through a Python backend that communicates with Unity using files or through a server-based communication protocol (like WebSockets or HTTP).

Here’s a simple Python script to parse the .pkl file, extract the keypoints, and output them in a way that can be transferred to Unity for further use.

Parser for .pkl Files Containing 3D Keypoints in COCO Format

import pickle
import numpy as np
import json

# Define the mapping for COCO keypoints if necessary
COCO_KEYPOINTS = [
    "nose", "left_eye", "right_eye", "left_ear", "right_ear",
    "left_shoulder", "right_shoulder", "left_elbow", "right_elbow",
    "left_wrist", "right_wrist", "left_hip", "right_hip",
    "left_knee", "right_knee", "left_ankle", "right_ankle"
]

def load_pkl_file(file_path):
    """Load a .pkl file and return its contents."""
    with open(file_path, 'rb') as f:
        data = pickle.load(f)
    return data

def extract_3d_keypoints(data):
    """Extract 3D keypoints from the loaded data."""
    keypoints_3d = data.get('keypoints_3d', None)
    if keypoints_3d is None:
        raise ValueError("No 'keypoints_3d' field found in the data.")

    return keypoints_3d

def convert_to_unity_format(keypoints_3d):
    """
    Convert the keypoints into a Unity-friendly format.
    Here we assume that the keypoints are in the form (x, y, z) for each point.
    We'll output it as a JSON string that can be passed to Unity.
    """
    keypoints_list = []

    for i, keypoint in enumerate(keypoints_3d):
        # Flattening 3D keypoints into a list for easy processing in Unity
        keypoint_dict = {
            "name": COCO_KEYPOINTS[i] if i < len(COCO_KEYPOINTS) else f"keypoint_{i}",
            "x": keypoint[0],
            "y": keypoint[1],
            "z": keypoint[2]
        }
        keypoints_list.append(keypoint_dict)

    return keypoints_list

def save_keypoints_to_json(keypoints, output_file):
    """Save the keypoints to a JSON file for Unity integration."""
    with open(output_file, 'w') as f:
        json.dump(keypoints, f, indent=4)

    print(f"Keypoints saved to {output_file}")

if __name__ == "__main__":
    # Example usage
    pkl_file_path = "path_to_your_pkl_file.pkl"
    output_json_path = "keypoints_unity.json"

    # Load .pkl file
    data = load_pkl_file(pkl_file_path)

    # Extract 3D keypoints from the data
    keypoints_3d = extract_3d_keypoints(data)

    # Convert keypoints to Unity-friendly format
    keypoints_unity_format = convert_to_unity_format(keypoints_3d)

    # Save the keypoints in a JSON file for Unity to use
    save_keypoints_to_json(keypoints_unity_format, output_json_path)

Key Aspects of the Code:

Loading .pkl File: The pickle module is used to load the .pkl file, which contains 3D keypoints.
Keypoint Extraction: The extract_3d_keypoints() function extracts the 3D keypoints from the data.
Unity-Compatible Format: The convert_to_unity_format() function converts the keypoints into a format that can easily be processed by Unity, e.g., a JSON format with a list of keypoints and their 3D coordinates.
JSON Export: The save_keypoints_to_json() function writes the extracted keypoints to a .json file, which Unity can then load and parse.

Unity Integration

On the Unity side, you can create a script to read the generated JSON file and process it into Unity’s 3D world.

Unity C# Script to Read Keypoints from JSON

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

public class KeypointLoader : MonoBehaviour
{
    [System.Serializable]
    public class Keypoint
    {
        public string name;
        public float x;
        public float y;
        public float z;
    }

    [System.Serializable]
    public class KeypointList
    {
        public List<Keypoint> keypoints;
    }

    public string filePath = "Assets/keypoints_unity.json";  // Path to the JSON file
    public GameObject[] keypointObjects;  // Reference to Unity GameObjects representing each keypoint

    void Start()
    {
        LoadKeypointsFromFile(filePath);
    }

    void LoadKeypointsFromFile(string path)
    {
        // Read the JSON file and parse the keypoints
        string jsonContent = File.ReadAllText(path);
        KeypointList keypointList = JsonUtility.FromJson<KeypointList>(jsonContent);

        // Assign positions to keypoint GameObjects in Unity
        for (int i = 0; i < keypointList.keypoints.Count; i++)
        {
            Keypoint keypoint = keypointList.keypoints[i];
            if (i < keypointObjects.Length)
            {
                Vector3 position = new Vector3(keypoint.x, keypoint.y, keypoint.z);
                keypointObjects[i].transform.position = position;
            }
        }
    }
}

Steps to Implement:

Parser: Run the Python parser script to extract and save 3D keypoints in JSON format.
Unity Setup:
- Create GameObjects to represent each keypoint.
- Link the keypointObjects array to the corresponding GameObjects.
- Load the JSON file at runtime in Unity and update the position of each GameObject based on the keypoint data.

This method should allow you to visualize or interact with the keypoints data in Unity.

To test the converter and visualize the 3D keypoints as an animated stick figure, we can use matplotlib to create a plot that shows the stick figure frame-by-frame. Since we’re working with 3D keypoints, we will use a 3D plot.

Below is the Python code to accomplish this. The stick figure will be animated and displayed frame-by-frame without saving images to disk.

Code to Display the Animated Stick Figure:

import pickle
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation

# Define the skeleton connections based on COCO format
COCO_CONNECTIONS = [
    (0, 1), (0, 2), (1, 3), (2, 4),  # Head
    (5, 6), (5, 7), (7, 9), (6, 8), (8, 10),  # Arms
    (5, 11), (6, 12), (11, 12),  # Torso
    (11, 13), (13, 15), (12, 14), (14, 16)  # Legs
]

def load_pkl_file(file_path):
    """Load a .pkl file and return its contents."""
    with open(file_path, 'rb') as f:
        data = pickle.load(f)
    return data

def extract_3d_keypoints(data):
    """Extract 3D keypoints from the loaded data."""
    keypoints_3d = data.get('keypoints_3d', None)
    if keypoints_3d is None:
        raise ValueError("No 'keypoints_3d' field found in the data.")

    return keypoints_3d

def plot_stick_figure(keypoints_3d, ax):
    """Plot the stick figure based on 3D keypoints."""
    ax.clear()

    # Plot keypoints
    xs, ys, zs = zip(*keypoints_3d)
    ax.scatter(xs, ys, zs, color='blue')

    # Plot skeleton connections
    for connection in COCO_CONNECTIONS:
        point1 = keypoints_3d[connection[0]]
        point2 = keypoints_3d[connection[1]]
        xs = [point1[0], point2[0]]
        ys = [point1[1], point2[1]]
        zs = [point1[2], point2[2]]
        ax.plot(xs, ys, zs, color='black')

    # Set axis limits for better visualization
    ax.set_xlim([-1, 1])
    ax.set_ylim([-1, 1])
    ax.set_zlim([0, 2])

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

def animate_stick_figure(frames):
    """Animate the stick figure frame-by-frame."""
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    def update(frame):
        plot_stick_figure(frames[frame], ax)

    anim = FuncAnimation(fig, update, frames=len(frames), interval=200)
    plt.show()

if __name__ == "__main__":
    # Example usage
    pkl_file_path = "path_to_your_pkl_file.pkl"

    # Load .pkl file
    data = load_pkl_file(pkl_file_path)

    # Extract 3D keypoints from the data (assuming it contains multiple frames)
    keypoints_3d_frames = extract_3d_keypoints(data)

    # Animate stick figure
    animate_stick_figure(keypoints_3d_frames)

Explanation of the Code:

Loading .pkl File: The .pkl file is loaded using the pickle module, and 3D keypoints are extracted.
Stick Figure Connections: The stick figure is created using the connections defined by the COCO format, which defines which keypoints are connected.
3D Plotting: Using matplotlib, we plot each frame in 3D, connecting the keypoints with lines to form a stick figure.
Animation: The FuncAnimation function from matplotlib is used to animate the stick figure, showing each frame of the keypoints in sequence.

Instructions to Use:

Replace "path_to_your_pkl_file.pkl" with the path to your .pkl file containing 3D keypoints in COCO format.
Run the code, and it will display the animated stick figure in a matplotlib window.

This code doesn’t save any images or frames to disk but will animate the stick figure in real-time using matplotlib’s interactive plot display.

To help you navigate the structure of your .pkl file, let’s print out the entire data hierarchy in a readable way. This will let us explore all available keys and the type of data contained within.

Here’s a Python script that recursively prints out the entire structure of your .pkl file, including the types of values associated with each key:

Code to Print the Data Hierarchy:

import pickle

def print_data_hierarchy(data, indent=0):
    """Recursively print the data hierarchy with indentation."""
    indent_str = "  " * indent
    if isinstance(data, dict):
        for key, value in data.items():
            print(f"{indent_str}{key}: {type(value)}")
            print_data_hierarchy(value, indent + 1)
    elif isinstance(data, list):
        print(f"{indent_str}List of {len(data)} items")
        if len(data) > 0:
            print_data_hierarchy(data[0], indent + 1)  # Print the structure of the first item
    elif isinstance(data, tuple):
        print(f"{indent_str}Tuple of {len(data)} items")
        for i, item in enumerate(data):
            print(f"{indent_str}Item {i}:")
            print_data_hierarchy(item, indent + 1)
    else:
        print(f"{indent_str}{type(data).__name__}")

def load_pkl_file(file_path):
    """Load a .pkl file and return its contents."""
    with open(file_path, 'rb') as f:
        data = pickle.load(f)
    return data

if __name__ == "__main__":
    # Load .pkl file
    pkl_file_path = "path_to_your_pkl_file.pkl"  # Replace with your file path
    data = load_pkl_file(pkl_file_path)

    # Print data hierarchy
    print("Data Hierarchy:")
    print_data_hierarchy(data)

Instructions:

Replace "path_to_your_pkl_file.pkl": Update this with the actual path to your .pkl file.
Run the script: This will print the entire hierarchy of your data.

Explanation:

The script uses recursion to walk through dictionaries, lists, and tuples to explore the nested structure.
It prints the type of each item to help you understand how the data is organized.

Output Example:

The output of the script will look something like this (hypothetical example):

Data Hierarchy:
keypoints: <class 'list'>
  List of 100 items
    List of 17 items
      <class 'tuple'>
      Tuple of 3 items
        <class 'float'>
        <class 'float'>
        <class 'float'>
video_metadata: <class 'dict'>
  video_id: <class 'str'>
  fps: <class 'float'>

This will allow you to identify where the relevant keypoints_3d or similar data is located in your .pkl file. Once you find the correct key or data structure, you can modify your extraction logic accordingly.

To zoom out the 3D plot so that the entire stick figure is visible, you can adjust the axis limits dynamically based on the range of keypoints. Specifically, you can calculate the min and max values for each axis (X, Y, and Z) and set the axis limits accordingly.

Here’s an updated version of the code that dynamically adjusts the axis limits to ensure the entire stick figure is visible:

Updated Code to Dynamically Set Axis Limits:

import pickle
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation

# Define the skeleton connections based on COCO format
COCO_CONNECTIONS = [
    (0, 1), (0, 2), (1, 3), (2, 4),  # Head
    (5, 6), (5, 7), (7, 9), (6, 8), (8, 10),  # Arms
    (5, 11), (6, 12), (11, 12),  # Torso
    (11, 13), (13, 15), (12, 14), (14, 16)  # Legs
]

def load_pkl_file(file_path):
    """Load a .pkl file and return its contents."""
    with open(file_path, 'rb') as f:
        data = pickle.load(f)
    return data

def extract_3d_keypoints(data):
    """Extract 3D keypoints from the loaded data."""
    keypoints_3d = data.get('keypoints_3d', None)
    if keypoints_3d is None:
        raise ValueError("No 'keypoints_3d' field found in the data.")

    return keypoints_3d

def plot_stick_figure(keypoints_3d, ax):
    """Plot the stick figure based on 3D keypoints."""
    ax.clear()

    # Plot keypoints
    xs, ys, zs = zip(*keypoints_3d)
    ax.scatter(xs, ys, zs, color='blue')

    # Plot skeleton connections
    for connection in COCO_CONNECTIONS:
        point1 = keypoints_3d[connection[0]]
        point2 = keypoints_3d[connection[1]]
        xs = [point1[0], point2[0]]
        ys = [point1[1], point2[1]]
        zs = [point1[2], point2[2]]
        ax.plot(xs, ys, zs, color='black')

    # Dynamically set axis limits to fit the stick figure
    ax.set_xlim([min(xs), max(xs)])
    ax.set_ylim([min(ys), max(ys)])
    ax.set_zlim([min(zs), max(zs)])

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

def animate_stick_figure(frames):
    """Animate the stick figure frame-by-frame."""
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    def update(frame):
        plot_stick_figure(frames[frame], ax)

    anim = FuncAnimation(fig, update, frames=len(frames), interval=200)
    plt.show()

if __name__ == "__main__":
    # Example usage
    pkl_file_path = "path_to_your_pkl_file.pkl"

    # Load .pkl file
    data = load_pkl_file(pkl_file_path)

    # Extract 3D keypoints from the data (assuming it contains multiple frames)
    keypoints_3d_frames = extract_3d_keypoints(data)

    # Animate stick figure
    animate_stick_figure(keypoints_3d_frames)

Changes Made:

Dynamic Axis Limits: The axis limits are now set dynamically based on the min and max values of the keypoints for each axis. This ensures that the stick figure is properly zoomed out and centered.
- ax.set_xlim([min(xs), max(xs)])
- ax.set_ylim([min(ys), max(ys)])
- ax.set_zlim([min(zs), max(zs)])
3D Viewport Adjustment: By dynamically setting the axis limits, the plot will automatically zoom out to fit the entire stick figure, regardless of the size or position of the figure.

Explanation:

The plot_stick_figure() function now calculates the minimum and maximum X, Y, and Z values from the keypoints and uses these to adjust the axis limits.
This ensures the figure is always properly framed and visible during the animation.

Try running this version of the code, and you should see the stick figure fully displayed, regardless of its size or position.

To hide the axis labels, ticks, and numbers while still showing the floor grid in a 3D plot, you can adjust the matplotlib axes configuration. We’ll modify the code to:

Hide axis labels.
Remove axis tick marks and numbers.
Keep the floor grid visible.

Here’s the updated code to achieve this:

Updated Code to Show Only the Floor Grid:

import pickle
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation

# Define the skeleton connections based on COCO format
COCO_CONNECTIONS = [
    (0, 1), (0, 2), (1, 3), (2, 4),  # Head
    (5, 6), (5, 7), (7, 9), (

Here is the relevant part of the matplotlib code to hide axis labels, ticks, and numbers while keeping the floor grid visible:

def plot_stick_figure(keypoints_3d, ax):
    """Plot the stick figure based on 3D keypoints."""
    ax.clear()

    # Plot keypoints
    xs, ys, zs = zip(*keypoints_3d)
    ax.scatter(xs, ys, zs, color='blue')

    # Plot skeleton connections
    for connection in COCO_CONNECTIONS:
        point1 = keypoints_3d[connection[0]]
        point2 = keypoints_3d[connection[1]]
        xs = [point1[0], point2[0]]
        ys = [point1[1], point2[1]]
        zs = [point1[2], point2[2]]
        ax.plot(xs, ys, zs, color='black')

    # Dynamically set axis limits to fit the stick figure
    ax.set_xlim([min(xs), max(xs)])
    ax.set_ylim([min(ys), max(ys)])
    ax.set_zlim([min(zs), max(zs)])

    # Hide the axis labels, ticks, and numbers
    ax.set_xticks([])  # Remove x-axis ticks
    ax.set_yticks([])  # Remove y-axis ticks
    ax.set_zticks([])  # Remove z-axis ticks
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.set_zlabel('')

    # Show the floor grid
    ax.grid(True)

Key Changes:

Remove axis ticks and labels: ax.set_xticks([]), ax.set_yticks([]), and ax.set_zticks([]) remove the ticks on all three axes.
Hide axis labels: ax.set_xlabel(''), ax.set_ylabel(''), and ax.set_zlabel('') remove the axis labels.
Keep the grid: ax.grid(True) ensures that the grid remains visible on the plot.

This configuration will display the 3D stick figure on a grid with no axis labels or ticks, providing a clean look focused on the figure and the floor grid.

To inverse the Y and Z axes in the plot, we can simply modify the axis limits by flipping the range of values. Here’s how you can incorporate this change along with hiding the ticks and labels:

Updated Code to Inverse Y and Z Axes, and Show Only the Floor Grid:

def plot_stick_figure(keypoints_3d, ax):
    """Plot the stick figure based on 3D keypoints."""
    ax.clear()

    # Plot keypoints
    xs, ys, zs = zip(*keypoints_3d)
    ax.scatter(xs, ys, zs, color='blue')

    # Plot skeleton connections
    for connection in COCO_CONNECTIONS:
        point1 = keypoints_3d[connection[0]]
        point2 = keypoints_3d[connection[1]]
        xs = [point1[0], point2[0]]
        ys = [point1[1], point2[1]]
        zs = [point1[2], point2[2]]
        ax.plot(xs, ys, zs, color='black')

    # Dynamically set axis limits to fit the stick figure
    ax.set_xlim([min(xs), max(xs)])

    # Inverse the Y-axis
    ax.set_ylim([max(ys), min(ys)])  # Flipping the Y-axis

    # Inverse the Z-axis
    ax.set_zlim([max(zs), min(zs)])  # Flipping the Z-axis

    # Hide the axis labels, ticks, and numbers
    ax.set_xticks([])  # Remove x-axis ticks
    ax.set_yticks([])  # Remove y-axis ticks
    ax.set_zticks([])  # Remove z-axis ticks
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.set_zlabel('')

    # Show the floor grid
    ax.grid(True)

Key Updates:

Inverse Y-axis: ax.set_ylim([max(ys), min(ys)]) flips the Y-axis.
Inverse Z-axis: ax.set_zlim([max(zs), min(zs)]) flips the Z-axis.
Grid Display: ax.grid(True) continues to display the floor grid.

This configuration will display the stick figure with the Y and Z axes inverted, as well as a clean plot with no axis labels, ticks, or numbers, while keeping the grid.

The error you’re encountering is due to the fact that NumPy arrays (ndarray) cannot be directly serialized into JSON format by Python’s json module. However, NumPy arrays can be easily converted to standard Python lists, which are serializable.

To resolve the issue, you can convert the NumPy arrays (e.g., keypoints data) to Python lists before serializing the data to JSON. Here’s how you can handle this:

Updated Code to Handle JSON Serialization of NumPy Arrays:

import json
import numpy as np
import pickle

def numpy_to_list(data):
    """Recursively convert numpy arrays to lists in a data structure."""
    if isinstance(data, np.ndarray):
        return data.tolist()  # Convert numpy array to list
    elif isinstance(data, dict):
        return {key: numpy_to_list(value) for key, value in data.items()}
    elif isinstance(data, list):
        return [numpy_to_list(item) for item in data]
    else:
        return data

def load_pkl_file(file_path):
    """Load a .pkl file and return its contents."""
    with open(file_path, 'rb') as f:
        data = pickle.load(f)
    return data

def save_keypoints_to_json(data, output_file):
    """Convert NumPy arrays to lists and save to a JSON file."""
    data_as_list = numpy_to_list(data)  # Convert any NumPy arrays to lists
    with open(output_file, 'w') as f:
        json.dump(data_as_list, f, indent=4)
    print(f"Data successfully saved to {output_file}")

if __name__ == "__main__":
    pkl_file_path = "path_to_your_pkl_file.pkl"
    output_json_path = "keypoints_data.json"

    # Load the .pkl file
    data = load_pkl_file(pkl_file_path)

    # Save the keypoints data to a JSON file after converting NumPy arrays to lists
    save_keypoints_to_json(data, output_json_path)

Explanation of the Changes:

numpy_to_list() Function: This function recursively walks through the data structure and converts any NumPy arrays (ndarray) to standard Python lists using .tolist().
save_keypoints_to_json() Function: Converts the NumPy arrays to lists and then saves the data in a JSON file using Python’s json.dump() method.
JSON Serialization: By converting the NumPy arrays to lists, Python’s json module will be able to handle the serialization without throwing the error.

How to Use:

Replace "path_to_your_pkl_file.pkl" with the actual path to your .pkl file.
Run the code, and the keypoints (or other data) will be serialized to JSON, with NumPy arrays automatically converted to lists.

This approach should resolve the TypeError: Object of type ndarray is not JSON serializable issue, allowing you to properly save the data to a JSON file.

this is wrong, because the data are organized in frames: def convert_to_unity_format(keypoints_3d): """ Convert the keypoints into a Unity-friendly format. Here we assume that the keypoints are in the form (x, y, z) for each point. We’ll output it as a JSON string that can be passed to Unity. """ keypoints_list = []

for i, keypoint in enumerate(keypoints_3d):
    # Flattening 3D keypoints into a list for easy processing in Unity
    keypoint_dict = {
        "name": COCO_KEYPOINTS[i] if i < len(COCO_KEYPOINTS) else f"keypoint_{i}",
        "x": keypoint[0],
        "y": keypoint[1],
        "z": keypoint[2]
    }
    keypoints_list.append(keypoint_dict)

return keypoints_list

You’re right! Since the data is organized in frames (i.e., each frame contains a list of keypoints), we need to update the convert_to_unity_format() function to properly handle the frame structure. In this case, each frame will have multiple keypoints, and we should format each frame separately while ensuring the data structure is friendly for Unity processing.

Updated Function to Handle Frame-by-Frame Keypoints

Here’s how you can modify the convert_to_unity_format() function to properly handle multiple frames of keypoints:

def convert_to_unity_format(keypoints_3d_frames):
    """
    Convert the keypoints into a Unity-friendly format.
    The keypoints_3d_frames is expected to be a list where each item is a frame,
    and each frame contains a list of keypoints in the form (x, y, z) for each point.
    We'll output a list of frames, where each frame contains a list of keypoints dictionaries.
    """
    all_frames_list = []

    # Loop through each frame of keypoints
    for frame in keypoints_3d_frames:
        frame_keypoints_list = []

        # Loop through each keypoint in the frame
        for i, keypoint in enumerate(frame):
            # Flattening 3D keypoints into a dictionary for easy processing in Unity
            keypoint_dict = {
                "name": COCO_KEYPOINTS[i] if i < len(COCO_KEYPOINTS) else f"keypoint_{i}",
                "x": keypoint[0],
                "y": keypoint[1],
                "z": keypoint[2]
            }
            frame_keypoints_list.append(keypoint_dict)

        # Append the keypoints for this frame to the list of all frames
        all_frames_list.append(frame_keypoints_list)

    return all_frames_list

Explanation of the Changes:

Frame-by-Frame Handling: The outer loop processes each frame of keypoints. Each frame is a list of 3D keypoints.
Per Keypoint Processing: For each keypoint in the frame, a dictionary is created with the name, x, y, and z values.
Returning a List of Frames: After processing all keypoints in a frame, we append that frame’s keypoints to all_frames_list. The final result is a list of frames, where each frame contains a list of keypoint dictionaries.

Example JSON Structure:

The output structure will be a list of frames, and each frame will contain a list of keypoint dictionaries. Here’s an example of what one frame might look like in the final output:

[
    [
        {"name": "nose", "x": 0.1, "y": 0.5, "z": 1.2},
        {"name": "left_eye", "x": 0.2, "y": 0.6, "z": 1.3},
        {"name": "right_eye", "x": 0.3, "y": 0.4, "z": 1.4},
        // ... other keypoints for this frame
    ],
    // ... more frames
]

How to Use:

The input keypoints_3d_frames should be a list of frames. Each frame should be a list of keypoints, where each keypoint is represented as a tuple or array of (x, y, z) coordinates.
Call this function to convert the 3D keypoints data into a Unity-friendly JSON format.

Once converted, you can serialize the output to JSON using json.dump() as we discussed earlier, ensuring proper serialization for Unity integration.

To efficiently parse and handle large datasets of keypoints in Unity, especially for real-time applications, you can use NativeArray in combination with structs. This approach minimizes memory overhead and leverages Unity’s Job System and Burst Compiler for performance.

Here’s a step-by-step guide to creating a parser for Unity using NativeArray and structs.

1. Define the `Keypoint` Struct

We will define a Keypoint struct to store the 3D coordinates of each keypoint. The Keypoint struct will be used in combination with NativeArray to store the keypoints for each frame.

using UnityEngine;
using Unity.Collections;

public struct Keypoint
{
    public Vector3 position;

    public Keypoint(float x, float y, float z)
    {
        position = new Vector3(x, y, z);
    }
}

2. Define the `KeypointFrame` Struct

Since keypoints are organized in frames, we also need a KeypointFrame struct to manage a set of keypoints for each frame. This struct will use a NativeArray<Keypoint>.

using Unity.Collections;
using UnityEngine;

public struct KeypointFrame
{
    public NativeArray<Keypoint> keypoints;

    public KeypointFrame(int numKeypoints, Allocator allocator)
    {
        keypoints = new NativeArray<Keypoint>(numKeypoints, allocator);
    }

    public void Dispose()
    {
        if (keypoints.IsCreated)
        {
            keypoints.Dispose();
        }
    }
}

3. Create a Parser to Convert JSON to NativeArrays

Now, let’s create the parser in Unity that converts the JSON data (from the previously serialized data) into NativeArray<Keypoint> structures. We will use the Unity.Collections.NativeArray for memory efficiency and performance.

Assuming you have a JSON file containing the keypoints for multiple frames, here’s the code that reads this JSON, processes it, and stores the keypoints in NativeArray:

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

[Serializable]
public class KeypointData
{
    public string name;
    public float x;
    public float y;
    public float z;
}

[Serializable]
public class FrameData
{
    public KeypointData[] keypoints;
}

[Serializable]
public class KeypointDataset
{
    public FrameData[] frames;
}

public class KeypointParser : MonoBehaviour
{
    public NativeArray<KeypointFrame> keypointFrames;
    private KeypointDataset dataset;

    // Path to the JSON file
    public string filePath = "Assets/keypoints_data.json";

    void Start()
    {
        // Load the dataset
        LoadKeypointDataset();
        // Allocate memory and fill the NativeArray with keypoint data
        ParseKeypoints();
    }

    void LoadKeypointDataset()
    {
        string json = File.ReadAllText(filePath);
        dataset = JsonUtility.FromJson<KeypointDataset>(json);
    }

    void ParseKeypoints()
    {
        // Allocate NativeArray for frames based on the number of frames in the dataset
        keypointFrames = new NativeArray<KeypointFrame>(dataset.frames.Length, Allocator.Persistent);

        for (int i = 0; i < dataset.frames.Length; i++)
        {
            FrameData frame = dataset.frames[i];

            // Create a NativeArray for the keypoints in this frame
            KeypointFrame keypointFrame = new KeypointFrame(frame.keypoints.Length, Allocator.Persistent);

            // Populate keypoints in the frame
            for (int j = 0; j < frame.keypoints.Length; j++)
            {
                KeypointData keypointData = frame.keypoints[j];
                keypointFrame.keypoints[j] = new Keypoint(keypointData.x, keypointData.y, keypointData.z);
            }

            // Assign the frame to the NativeArray
            keypointFrames[i] = keypointFrame;
        }
    }

    void OnDestroy()
    {
        // Dispose of NativeArray memory when no longer needed
        if (keypointFrames.IsCreated)
        {
            for (int i = 0; i < keypointFrames.Length; i++)
            {
                keypointFrames[i].Dispose();
            }
            keypointFrames.Dispose();
        }
    }
}

Breakdown of the Parser:

Struct Definitions:
- Keypoint: Holds the Vector3 for each keypoint’s position.
- KeypointFrame: Holds a NativeArray<Keypoint> for keypoints in a single frame and handles memory allocation/disposal.
Loading JSON:
- The JSON is loaded using File.ReadAllText(), and then parsed with JsonUtility.FromJson<KeypointDataset>() into an object that can be easily processed.
Memory Management:
- For each frame, a NativeArray<Keypoint> is created using the Allocator.Persistent allocator, which ensures that the memory will remain allocated until manually disposed of.
Populating the Data:
- The keypoint data from JSON is populated into NativeArray<Keypoint> inside the KeypointFrame.
Cleanup:
- Since NativeArray requires manual memory management, we use the OnDestroy method to dispose of the arrays when the object is destroyed.

Example JSON Data Structure:

Your JSON data should follow this structure to match the parsing logic:

{
  "frames": [
    {
      "keypoints": [
        {"name": "nose", "x": 0.1, "y": 0.5, "z": 1.2},
        {"name": "left_eye", "x": 0.2, "y": 0.6, "z": 1.3},
        {"name": "right_eye", "x": 0.3, "y": 0.4, "z": 1.4}
        // ... other keypoints
      ]
    },
    {
      "keypoints": [
        {"name": "nose", "x": 0.2, "y": 0.4, "z": 1.0},
        {"name": "left_eye", "x": 0.3, "y": 0.5, "z": 1.1},
        {"name": "right_eye", "x": 0.4, "y": 0.3, "z": 1.2}
        // ... other keypoints
      ]
    }
    // ... more frames
  ]
}

Why Use `NativeArray`?

Performance: NativeArray allows for efficient memory allocation and is compatible with Unity’s Job System, enabling parallel processing.
Manual Memory Management: With NativeArray, you explicitly control the memory, which is crucial for performance-sensitive applications like real-time keypoint processing in 3D animations.

How to Use:

Add the script to a Unity GameObject.
Make sure the JSON file path is correct, and the data is in the right format.
Run the Unity project, and it will load the keypoints, store them in NativeArray, and allow further manipulation or visualization.

This setup is efficient for managing large datasets and ensures that your application can process the keypoints in real-time using Unity’s performance optimizations.

To build a highly efficient, data-oriented system for loading and saving .pkl files, processing frames, and exporting the data in a format that Unity can directly read and convert to a NativeArray<Vector3>, we need to focus on minimizing memory allocations and classes.

We will:

Read all .pkl files from a folder.
Extract the 17 joints (51 floats per frame).
Create a header with the number of .pkl files and the frame count per .pkl.
Write all data in a single binary file with minimal overhead.
Deserialize this data in Unity using NativeArray<float> and reinterpret it as Vector3.

Python: Writing the Binary File

Here’s a Python script that reads the .pkl files, extracts the frame data, and writes it to a binary file in the desired format.

import os
import pickle
import struct
import numpy as np

# Path to the folder containing .pkl files
folder_path = 'path_to_your_pkl_folder'
output_file = 'output_keypoints_data.bin'

# Define the number of joints (17 joints with 3 coordinates each)
JOINT_COUNT = 17
COORDS_PER_JOINT = 3
FLOATS_PER_FRAME = JOINT_COUNT * COORDS_PER_JOINT

def load_pkl_files(folder_path):
    """Load all .pkl files from a folder and return the data and frame counts."""
    pkl_files = [f for f in os.listdir(folder_path) if f.endswith('.pkl')]
    frame_data_list = []
    frames_per_pkl = []

    for pkl_file in pkl_files:
        file_path = os.path.join(folder_path, pkl_file)

        with open(file_path, 'rb') as f:
            data = pickle.load(f)

        # Extract 3D keypoints from the data (17 joints, 3 coordinates per joint per frame)
        keypoints_3d = np.array(data.get('keypoints_3d', []))

        # Reshape keypoints into flat array of floats (17 * 3 floats per frame)
        if keypoints_3d.shape[1:] != (JOINT_COUNT, COORDS_PER_JOINT):
            raise ValueError(f"Unexpected keypoint shape in file {pkl_file}")

        # Flatten frames into a 1D list of floats
        flattened_frames = keypoints_3d.reshape(-1)
        frame_data_list.append(flattened_frames)
        frames_per_pkl.append(len(keypoints_3d))  # Store frame count for each file

    return frame_data_list, frames_per_pkl, len(pkl_files)

def write_binary_file(output_file, frame_data_list, frames_per_pkl, pkl_count):
    """Write all data to a single binary file."""
    with open(output_file, 'wb') as f:
        # Write the .pkl count (first line)
        f.write(struct.pack('I', pkl_count))  # 'I' = unsigned int (4 bytes)

        # Write the frames per .pkl list (second line)
        f.write(struct.pack(f'{pkl_count}I', *frames_per_pkl))  # List of frame counts

        # Write all the keypoint data as floats
        for frame_data in frame_data_list:
            f.write(struct.pack(f'{len(frame_data)}f', *frame_data))  # All floats

if __name__ == "__main__":
    # Load all .pkl files and gather the keypoints data
    frame_data_list, frames_per_pkl, pkl_count = load_pkl_files(folder_path)

    # Write the data to a binary file
    write_binary_file(output_file, frame_data_list, frames_per_pkl, pkl_count)

    print(f"Data successfully written to {output_file}")

Breakdown:

load_pkl_files():
- Loads all .pkl files from the folder.
- Extracts the keypoints_3d field from each .pkl file (assuming it’s organized as [frames, 17, 3]).
- Flattens the keypoints into a list of floats and stores them.
write_binary_file():
- First writes the number of .pkl files (pkl_count).
- Then writes the number of frames in each .pkl file.
- Finally, writes all the frame data as a continuous array of floats.

Binary File Structure:

First 4 bytes: Unsigned integer representing the number of .pkl files.
Next 4 bytes × number of .pkl files: Frame count for each .pkl file.
Remaining data: Flattened frame data as 32-bit floats.

Unity: Reading and Deserializing the Data

Here’s the Unity C# code to read this binary file, deserialize the data, and load it into a NativeArray<float>, which will then be reinterpreted as Vector3 data for use in Unity.

using System;
using System.IO;
using Unity.Collections;
using UnityEngine;

public class KeypointLoader : MonoBehaviour
{
    private NativeArray<Vector3> keypointData;
    private int pklCount;
    private int[] framesPerPkl;

    // Path to the binary file
    public string filePath = "Assets/output_keypoints_data.bin";

    void Start()
    {
        LoadKeypointData(filePath);
    }

    void LoadKeypointData(string path)
    {
        // Read the binary file into a byte array
        byte[] data = File.ReadAllBytes(path);

        // Create a binary reader for the data
        using (BinaryReader reader = new BinaryReader(new MemoryStream(data)))
        {
            // Read .pkl count (unsigned int, 4 bytes)
            pklCount = reader.ReadInt32();

            // Read frames per .pkl (each is an unsigned int, 4 bytes per frame count)
            framesPerPkl = new int[pklCount];
            for (int i = 0; i < pklCount; i++)
            {
                framesPerPkl[i] = reader.ReadInt32();
            }

            // Calculate total frame count and data size
            int totalFrames = 0;
            foreach (int frameCount in framesPerPkl)
            {
                totalFrames += frameCount;
            }

            // Total number of floats (each frame has 17 joints * 3 coords = 51 floats)
            int totalFloats = totalFrames * 51;

            // Create a NativeArray to hold the data
            NativeArray<float> floatData = new NativeArray<float>(totalFloats, Allocator.Persistent);

            // Read all float data from the binary file
            for (int i = 0; i < totalFloats; i++)
            {
                floatData[i] = reader.ReadSingle();
            }

            // Reinterpret the NativeArray of floats as a NativeArray of Vector3
            keypointData = floatData.Reinterpret<Vector3>(sizeof(float) * 3);

            Debug.Log($"Successfully loaded {pklCount} .pkl files with a total of {totalFrames} frames.");
        }
    }

    void OnDestroy()
    {
        // Dispose of the NativeArray when done
        if (keypointData.IsCreated)
        {
            keypointData.Dispose();
        }
    }
}

Breakdown:

Binary File Reading: The binary file is read using a BinaryReader. First, the .pkl count and the frame counts are read, then the entire float data is loaded.
NativeArray<float> to NativeArray<Vector3>: The float array is reinterpreted as Vector3 using NativeArray.Reinterpret.
Memory Management: The NativeArray is allocated using Allocator.Persistent and disposed of in the OnDestroy() method to avoid memory leaks.

Key Points for Efficiency:

No Intermediate Classes: We directly load the data into a NativeArray, avoiding unnecessary allocations or classes.
Direct Memory Reinterpretation: The float data is reinterpreted as Vector3 without additional copying or transformations.
Minimal Allocations: Data is stored in binary format for minimal overhead, and memory is managed manually using NativeArray.

This approach ensures efficient loading and memory usage, making it highly suitable for real-time applications such as animations or visualizations in Unity.

To integrate the tqdm library into the Python script to report the progress of .pkl parsing, we can simply wrap the iteration over the .pkl files with tqdm. This will give us a progress bar that updates as we process each file.

Below is the updated Python script that uses tqdm to show progress during the loading and parsing of .pkl files.

Python Script with `tqdm` for Progress Reporting

import os
import pickle
import struct
import numpy as np
from tqdm import tqdm  # Import tqdm for progress reporting

# Path to the folder containing .pkl files
folder_path = 'path_to_your_pkl_folder'
output_file = 'output_keypoints_data.bin'

# Define the number of joints (17 joints with 3 coordinates each)
JOINT_COUNT = 17
COORDS_PER_JOINT = 3
FLOATS_PER_FRAME = JOINT_COUNT * COORDS_PER_JOINT

def load_pkl_files(folder_path):
    """Load all .pkl files from a folder and return the data and frame counts."""
    pkl_files = [f for f in os.listdir(folder_path) if f.endswith('.pkl')]
    frame_data_list = []
    frames_per_pkl = []

    # Use tqdm to report progress as we iterate over pkl files
    for pkl_file in tqdm(pkl_files, desc="Parsing .pkl files"):
        file_path = os.path.join(folder_path, pkl_file)

        with open(file_path, 'rb') as f:
            data = pickle.load(f)

        # Extract 3D keypoints from the data (17 joints, 3 coordinates per joint per frame)
        keypoints_3d = np.array(data.get('keypoints_3d', []))

        # Reshape keypoints into flat array of floats (17 * 3 floats per frame)
        if keypoints_3d.shape[1:] != (JOINT_COUNT, COORDS_PER_JOINT):
            raise ValueError(f"Unexpected keypoint shape in file {pkl_file}")

        # Flatten frames into a 1D list of floats
        flattened_frames = keypoints_3d.reshape(-1)
        frame_data_list.append(flattened_frames)
        frames_per_pkl.append(len(keypoints_3d))  # Store frame count for each file

    return frame_data_list, frames_per_pkl, len(pkl_files)

def write_binary_file(output_file, frame_data_list, frames_per_pkl, pkl_count):
    """Write all data to a single binary file."""
    with open(output_file, 'wb') as f:
        # Write the .pkl count (first line)
        f.write(struct.pack('I', pkl_count))  # 'I' = unsigned int (4 bytes)

        # Write the frames per .pkl list (second line)
        f.write(struct.pack(f'{pkl_count}I', *frames_per_pkl))  # List of frame counts

        # Write all the keypoint data as floats
        for frame_data in tqdm(frame_data_list, desc="Writing frames to binary file"):
            f.write(struct.pack(f'{len(frame_data)}f', *frame_data))  # All floats

if __name__ == "__main__":
    # Load all .pkl files and gather the keypoints data
    frame_data_list, frames_per_pkl, pkl_count = load_pkl_files(folder_path)

    # Write the data to a binary file
    write_binary_file(output_file, frame_data_list, frames_per_pkl, pkl_count)

    print(f"Data successfully written to {output_file}")

Changes:

tqdm Integration:
- tqdm is used in the for pkl_file in tqdm(...) loop to display progress while parsing .pkl files.
- A second tqdm progress bar is used during the binary file writing process in write_binary_file() to show progress as frames are written to the file.
Progress Bars:
- The progress bars will display the number of .pkl files parsed and the number of frames written.
- The desc argument of tqdm provides a label for each progress bar.

Output:

When running this script, you will see two progress bars: one for parsing the .pkl files and another for writing the binary data.

Example of progress bar:

Parsing .pkl files: 100%|█████████████████████████| 50/50 [00:10<00:00, 4.80it/s]
Writing frames to binary file: 100%|████████████████████| 500/500 [00:02<00:00, 200.00it/s]

This will give you real-time feedback on how much of the work is done and how long it might take to complete. It’s a simple but effective way to monitor progress, especially when processing large datasets with many .pkl files.

instead of this slow reading, use pointers to directly copy the memory to the nativearray

    NativeArray<float> floatData = new NativeArray<float>(totalFloats, Allocator.Persistent);

    // Read all float data from the binary file
    for (int i = 0; i < totalFloats; i++)
    {
        floatData[i] = reader.ReadSingle();
    }

To improve the performance of reading float data into a NativeArray<float> by directly copying memory, we can use pointers and unsafe code in Unity. This avoids reading each float one by one and instead performs a bulk memory copy from the byte buffer directly into the NativeArray<float>.

Here’s how to efficiently read the float data from a binary file and directly copy it into the NativeArray<float> using pointers in Unity.

C# Code with Direct Memory Copy Using Pointers

using System;
using System.IO;
using Unity.Collections;
using UnityEngine;
using System.Runtime.InteropServices; // For Marshal

public class KeypointLoader : MonoBehaviour
{
    private NativeArray<Vector3> keypointData;
    private int pklCount;
    private int[] framesPerPkl;

    // Path to the binary file
    public string filePath = "Assets/output_keypoints_data.bin";

    void Start()
    {
        LoadKeypointData(filePath);
    }

    void LoadKeypointData(string path)
    {
        // Read the binary file into a byte array
        byte[] data = File.ReadAllBytes(path);

        // Create a binary reader for the data
        using (BinaryReader reader = new BinaryReader(new MemoryStream(data)))
        {
            // Read .pkl count (unsigned int, 4 bytes)
            pklCount = reader.ReadInt32();

            // Read frames per .pkl (each is an unsigned int, 4 bytes per frame count)
            framesPerPkl = new int[pklCount];
            for (int i = 0; i < pklCount; i++)
            {
                framesPerPkl[i] = reader.ReadInt32();
            }

            // Calculate total frame count and data size
            int totalFrames = 0;
            foreach (int frameCount in framesPerPkl)
            {
                totalFrames += frameCount;
            }

            // Total number of floats (each frame has 17 joints * 3 coords = 51 floats)
            int totalFloats = totalFrames * 51;

            // Create a NativeArray to hold the data
            NativeArray<float> floatData = new NativeArray<float>(totalFloats, Allocator.Persistent);

            // Read all float data at once and directly copy to NativeArray using unsafe memory copy
            ReadFloatsDirectlyIntoNativeArray(reader, floatData, totalFloats);

            // Reinterpret the NativeArray of floats as a NativeArray of Vector3
            keypointData = floatData.Reinterpret<Vector3>(sizeof(float) * 3);

            Debug.Log($"Successfully loaded {pklCount} .pkl files with a total of {totalFrames} frames.");
        }
    }

    unsafe void ReadFloatsDirectlyIntoNativeArray(BinaryReader reader, NativeArray<float> floatData, int totalFloats)
    {
        // Get the pointer to the internal memory of the NativeArray
        IntPtr nativeArrayPtr = (IntPtr)NativeArrayUnsafeUtility.GetUnsafePtr(floatData);

        // Allocate a byte array to read all float data from the file (4 bytes per float)
        byte[] floatBytes = reader.ReadBytes(totalFloats * sizeof(float));

        // Use Marshal.Copy to copy the byte array directly into the NativeArray memory
        Marshal.Copy(floatBytes, 0, nativeArrayPtr, floatBytes.Length);
    }

    void OnDestroy()
    {
        // Dispose of the NativeArray when done
        if (keypointData.IsCreated)
        {
            keypointData.Dispose();
        }
    }
}

Key Changes:

Direct Memory Copy Using Pointers:
- NativeArrayUnsafeUtility.GetUnsafePtr(): This method gives us direct access to the memory pointer of the NativeArray<float>.
- Marshal.Copy(): This function is used to copy the byte array directly into the NativeArray memory in one go, instead of iterating over each float.
Unsafe Code:
- The ReadFloatsDirectlyIntoNativeArray method is marked as unsafe to allow direct memory access. This allows faster copying by leveraging pointers and minimizing unnecessary processing.
Binary Reading:
- Instead of reading each float individually, we first read the entire float data into a byte array (floatBytes), then copy it directly into the NativeArray using Marshal.Copy.

Explanation:

Reading the Binary Data: We read the entire float data at once into a byte array (floatBytes) instead of reading each float one by one.
Direct Memory Copy: By using the Marshal.Copy function, we directly copy the byte array into the NativeArray<float>’s memory, which is much faster than individual float reads.

Performance Benefits:

Batch Memory Transfer: Reading the data in one go and performing a single memory copy minimizes the time spent on I/O and memory allocation.
Reduced Loop Overhead: Since we are not looping through each float, we eliminate the overhead caused by multiple calls to reader.ReadSingle().

This approach ensures that your float data is loaded into Unity with minimal overhead, making it suitable for real-time applications or large datasets where performance is critical.

You’re right! Since the BinaryReader already works on a MemoryStream that contains the binary data in memory, there is no need to read the data twice. We can directly copy the data from the MemoryStream into the NativeArray without allocating another byte array.

Here’s an improved version of the code, avoiding redundant reading and directly copying from the MemoryStream to the NativeArray using the Stream.Read() method and Marshal.Copy.

Updated C# Code for Efficient Memory Copy:

using System;
using System.IO;
using Unity.Collections;
using UnityEngine;
using System.Runtime.InteropServices; // For Marshal

public class KeypointLoader : MonoBehaviour
{
    private NativeArray<Vector3> keypointData;
    private int pklCount;
    private int[] framesPerPkl;

    // Path to the binary file
    public string filePath = "Assets/output_keypoints_data.bin";

    void Start()
    {
        LoadKeypointData(filePath);
    }

    void LoadKeypointData(string path)
    {
        // Read the binary file into a byte array
        byte[] data = File.ReadAllBytes(path);

        // Create a binary reader for the data
        using (MemoryStream memoryStream = new MemoryStream(data))
        using (BinaryReader reader = new BinaryReader(memoryStream))
        {
            // Read .pkl count (unsigned int, 4 bytes)
            pklCount = reader.ReadInt32();

            // Read frames per .pkl (each is an unsigned int, 4 bytes per frame count)
            framesPerPkl = new int[pklCount];
            for (int i = 0; i < pklCount; i++)
            {
                framesPerPkl[i] = reader.ReadInt32();
            }

            // Calculate total frame count and data size
            int totalFrames = 0;
            foreach (int frameCount in framesPerPkl)
            {
                totalFrames += frameCount;
            }

            // Total number of floats (each frame has 17 joints * 3 coords = 51 floats)
            int totalFloats = totalFrames * 51;

            // Create a NativeArray to hold the data
            NativeArray<float> floatData = new NativeArray<float>(totalFloats, Allocator.Persistent);

            // Directly copy the float data from the MemoryStream into the NativeArray
            ReadFloatsDirectlyIntoNativeArray(memoryStream, floatData, totalFloats);

            // Reinterpret the NativeArray of floats as a NativeArray of Vector3
            keypointData = floatData.Reinterpret<Vector3>(sizeof(float) * 3);

            Debug.Log($"Successfully loaded {pklCount} .pkl files with a total of {totalFrames} frames.");
        }
    }

    unsafe void ReadFloatsDirectlyIntoNativeArray(MemoryStream stream, NativeArray<float> floatData, int totalFloats)
    {
        // Get the pointer to the internal memory of the NativeArray
        IntPtr nativeArrayPtr = (IntPtr)NativeArrayUnsafeUtility.GetUnsafePtr(floatData);

        // Directly copy the data from the MemoryStream into the NativeArray
        byte[] buffer = new byte[totalFloats * sizeof(float)];
        stream.Read(buffer, 0, buffer.Length);  // Read from current stream position

        // Copy from the buffer to the NativeArray memory using Marshal.Copy
        Marshal.Copy(buffer, 0, nativeArrayPtr, buffer.Length);
    }

    void OnDestroy()
    {
        // Dispose of the NativeArray when done
        if (keypointData.IsCreated)
        {
            keypointData.Dispose();
        }
    }
}

Key Improvements:

No Redundant Reading:
- We avoid the unnecessary second read by directly copying data from the MemoryStream, which already holds the entire file in memory.
Direct Copy from MemoryStream:
- Instead of reading into an intermediate byte array and then copying, we now directly read the needed data from the MemoryStream in chunks using Stream.Read().
Avoiding a Second Byte Array:
- The Stream.Read() method pulls the required data from the MemoryStream into a single buffer that’s immediately copied into the NativeArray. This reduces overhead.

Explanation of Changes:

MemoryStream: Since BinaryReader is reading from a MemoryStream, which already contains the entire file, there’s no need for a second pass to load the data. We can directly use MemoryStream.Read() to access the float data.
Direct Data Copy: The key function ReadFloatsDirectlyIntoNativeArray() now reads bytes directly from the MemoryStream, avoiding any unnecessary intermediate allocations or transformations.

Performance Benefits:

Single Pass Data Copy: By reading directly from the MemoryStream and using Marshal.Copy, this approach eliminates the need to allocate an additional array or perform multiple reads.
Efficient Memory Handling: The use of pointers (NativeArrayUnsafeUtility.GetUnsafePtr()) and Marshal.Copy ensures that the memory copy is done efficiently and in a single operation.

Usage:

Place this script in your Unity project.
Ensure the path to the binary file is correctly set in filePath.
Run the project, and the data will be loaded into a NativeArray<Vector3>, ready for use in Unity.

This version minimizes the overhead and memory allocations, making it ideal for efficiently handling large datasets, such as multiple .pkl files containing many frames of keypoint data.

Ha! I wrote a much better version than you.

Now check for any remainining redundancies

    var pointer = 0;

    // Read the binary file into a byte array
    var data = await File.ReadAllBytesAsync(path);

    // Create a binary reader for the data
    using var reader = new BinaryReader(new MemoryStream(data));

    // Read .pkl count (unsigned int, 4 bytes)
    pklCount = reader.ReadInt32();
    pointer+=4;

    // Read frames per .pkl (each is an unsigned int, 4 bytes per frame count)
    framesPerPkl = new int[pklCount];
    for (int i = 0; i < pklCount; i++)
    {
        framesPerPkl[i] = reader.ReadInt32();
        pointer+=4;
    }

    // Calculate total frame count and data size
    var totalFrames = framesPerPkl.Sum();

    // Total number of floats (each frame has 17 joints * 3 coords = 51 floats)
    int totalFloats = totalFrames * 51;

    var mem = new Memory<byte>(data, pointer, data.Length-pointer);

    if (keypointData.IsCreated)
        keypointData.Dispose();

    // Create a NativeArray to hold the data
    keypointData = new NativeArray<Vector3>(totalFloats/3, Allocator.Persistent);

    unsafe
    {
        UnsafeUtility.MemMove(keypointData.GetUnsafePtr(), mem.Pin().Pointer, data.Length-pointer);
    }

    Debug.Log($"Successfully loaded {pklCount} .pkl files with a total of {totalFrames} frames.");

Your approach is already quite efficient, using memory manipulation and direct pointer operations. However, there are still a couple of small optimizations and improvements that could be made to further streamline the code:

Redundancies or Improvements:

Unnecessary BinaryReader:
- You’re using a BinaryReader, but since you’re directly manipulating the byte array (data) with a Memory<byte>, you could eliminate the BinaryReader altogether and directly parse the byte data. This avoids the overhead of using BinaryReader and keeping track of pointer manually for both the reader and the byte array at the same time.
Memory Span Instead of Memory<byte>:
- Instead of using Memory<byte>, you can use Span<byte>, which offers better performance because it is not meant for async use and has a slightly lower overhead than Memory<byte> for scenarios where async is not needed. This is especially relevant since you’re working with raw memory in an unsafe block. Using Span<byte> would allow you to manipulate the data more efficiently.
Avoid Pinning Twice:
- When you’re calling mem.Pin().Pointer, you’re pinning the memory for the operation. While this is fine, it could be further optimized by directly working with the data in the original array rather than extracting a slice (mem) and then pinning it again.
Check for Disposal Earlier:
- Currently, you’re checking if keypointData.IsCreated after parsing the data. For safety and proper memory management, it’s better to check and dispose of the existing NativeArray at the very beginning of the function. This way, in the case of an exception occurring in the middle of execution, the data is still managed properly.

Suggested Optimized Code:

using System;
using System.IO;
using System.Linq;
using System.Runtime.CompilerServices;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using UnityEngine;

public class KeypointLoader : MonoBehaviour
{
    private NativeArray<Vector3> keypointData;
    private int pklCount;
    private int[] framesPerPkl;

    public async void LoadKeypointData(string path)
    {
        // Dispose existing NativeArray if necessary
        if (keypointData.IsCreated)
            keypointData.Dispose();

        // Read the binary file into a byte array
        var data = await File.ReadAllBytesAsync(path);

        // Read .pkl count (unsigned int, 4 bytes)
        pklCount = BitConverter.ToInt32(data, 0);

        // Read frames per .pkl (each is an unsigned int, 4 bytes per frame count)
        framesPerPkl = new int[pklCount];
        int pointer = 4;
        for (int i = 0; i < pklCount; i++)
        {
            framesPerPkl[i] = BitConverter.ToInt32(data, pointer);
            pointer += 4;
        }

        // Calculate total frame count and data size
        var totalFrames = framesPerPkl.Sum();
        int totalFloats = totalFrames * 51;  // 51 floats per frame

        // Allocate NativeArray to hold all the Vector3 data (floats / 3)
        keypointData = new NativeArray<Vector3>(totalFloats / 3, Allocator.Persistent);

        // Use Span<byte> for more efficient memory slicing
        var span = new Span<byte>(data, pointer, data.Length - pointer);

        unsafe
        {
            // Use UnsafeUtility.MemCpy instead of MemMove (since there's no overlap in memory here)
            UnsafeUtility.MemCpy(keypointData.GetUnsafePtr(), Unsafe.AsPointer(ref span[0]), span.Length);
        }

        Debug.Log($"Successfully loaded {pklCount} .pkl files with a total of {totalFrames} frames.");
    }

    private void OnDestroy()
    {
        // Dispose of the NativeArray to avoid memory leaks
        if (keypointData.IsCreated)
        {
            keypointData.Dispose();
        }
    }
}

Key Improvements:

Removed BinaryReader: Instead of using BinaryReader, the code directly reads the byte array (data) using BitConverter to convert the bytes into integers. This simplifies the logic by keeping everything in the raw byte domain without needing an additional stream reader.
Span<byte> Instead of Memory<byte>: We now use Span<byte> to work directly with the memory in a synchronous context, which eliminates the overhead of managing async-safe types like Memory<byte>.
Used UnsafeUtility.MemCpy: UnsafeUtility.MemCpy is used instead of MemMove because there’s no overlapping memory here. MemCpy is slightly more efficient for copying non-overlapping blocks of memory.
Check for Disposal Early: The disposal check for keypointData is moved to the beginning of the method to ensure proper memory management before starting any operations.

Performance Improvements:

Direct Memory Access: Using Span<byte> and eliminating unnecessary abstractions (like BinaryReader) allows you to directly access the memory in a faster and more efficient way.
Reduced Memory Allocations: By avoiding intermediate streams and abstractions, and using direct memory copying, the overall memory footprint and execution time are reduced.
Faster Copy: Using UnsafeUtility.MemCpy speeds up the memory copy operation because there’s no risk of overlap, and it’s optimized for such use cases.

This should now be the most efficient way to load your data with the least amount of overhead and redundancy.

You’re absolutely right — Span<T> is a stack-only type and cannot be used across async boundaries because it doesn’t persist across async/await calls. This is a limitation of Span<T>, as it is designed for performance by avoiding heap allocation, but it doesn’t work well in async methods.

Let’s fix the issue by using Memory<T> instead of Span<T>, as Memory<T> can be used with async methods and avoids the stack-only limitation.

Here’s how to update the code to handle the Memory<byte> correctly, ensuring async compatibility while still using direct memory copying:

Corrected Code Using `Memory<byte>`:

using System;
using System.IO;
using System.Linq;
using System.Runtime.CompilerServices;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using UnityEngine;

public class KeypointLoader : MonoBehaviour
{
    private NativeArray<Vector3> keypointData;
    private int pklCount;
    private int[] framesPerPkl;

    public async void LoadKeypointData(string path)
    {
        // Dispose existing NativeArray if necessary
        if (keypointData.IsCreated)
            keypointData.Dispose();

        // Read the binary file into a byte array asynchronously
        var data = await File.ReadAllBytesAsync(path);

        // Read .pkl count (unsigned int, 4 bytes)
        pklCount = BitConverter.ToInt32(data, 0);

        // Read frames per .pkl (each is an unsigned int, 4 bytes per frame count)
        framesPerPkl = new int[pklCount];
        int pointer = 4;
        for (int i = 0; i < pklCount; i++)
        {
            framesPerPkl[i] = BitConverter.ToInt32(data, pointer);
            pointer += 4;
        }

        // Calculate total frame count and data size
        var totalFrames = framesPerPkl.Sum();
        int totalFloats = totalFrames * 51;  // 51 floats per frame

        // Allocate NativeArray to hold all the Vector3 data (floats / 3)
        keypointData = new NativeArray<Vector3>(totalFloats / 3, Allocator.Persistent);

        // Create Memory<byte> to represent the float data slice starting at `pointer`
        var memory = new Memory<byte>(data, pointer, data.Length - pointer);

        unsafe
        {
            // Pin the memory and copy it directly into the NativeArray
            using (var handle = memory.Pin())
            {
                UnsafeUtility.MemCpy(keypointData.GetUnsafePtr(), (void*)handle.Pointer, memory.Length);
            }
        }

        Debug.Log($"Successfully loaded {pklCount} .pkl files with a total of {totalFrames} frames.");
    }

    private void OnDestroy()
    {
        // Dispose of the NativeArray to avoid memory leaks
        if (keypointData.IsCreated)
        {
            keypointData.Dispose();
        }
    }
}

Explanation of Changes:

Using Memory<byte>:
- Memory<byte> is now used instead of Span<byte> because it supports async methods and can be safely used across async/await boundaries. It also avoids stack allocation issues.
Pinning Memory<byte>:
- We use Memory<byte>.Pin() to pin the memory so we can access the underlying pointer, which is required for direct memory copying. The Pin method ensures the GC doesn’t move the data during the copy.
- Pin returns a MemoryHandle, which must be disposed after use (hence the using block).
Copying Data:
- Once the memory is pinned, we use UnsafeUtility.MemCpy() to efficiently copy the data from the pinned memory to the NativeArray<float>, using pointers.
Async Safety:
- This method is fully compatible with asynchronous patterns. Since Memory<byte> can exist across async methods, it works well with await File.ReadAllBytesAsync().

Performance Benefits:

Asynchronous File Reading: The file reading remains asynchronous, allowing the main thread to remain responsive while loading the binary data.
Direct Memory Copy: The memory is copied efficiently using UnsafeUtility.MemCpy(), ensuring minimal overhead during the transfer of data into NativeArray<Vector3>.
No Redundant Allocations: By avoiding intermediate structures and using the built-in memory utilities (Memory<byte> and Pin), memory allocations are minimized.

Usage:

Place this script in your Unity project.
Ensure the path to the binary file is correctly set in filePath.
Run the project, and the data will be asynchronously loaded into a NativeArray<Vector3>, ready for use in Unity.

This corrected version ensures both async compatibility and efficient memory handling, making it ideal for handling large datasets or multiple .pkl files.

can’t we avoid the loop here and memcopy the array directly?

    // Read frames per .pkl (each is an unsigned int, 4 bytes per frame count)
    framesPerPkl = new int[pklCount];
    for (var i = 0; i < pklCount; i++)
    {
        framesPerPkl[i] = BitConverter.ToInt32(data, pointer);
        pointer+=4;
    }

Yes, we can avoid the loop when reading the framesPerPkl array by directly copying the memory from the data byte array into the framesPerPkl array using Buffer.BlockCopy or a similar memory-copying mechanism. Since framesPerPkl is an array of integers (4 bytes each), we can copy the entire block of memory for the frame counts in one go, rather than iterating through them one by one.

Here’s how you can modify the code to directly copy the data for framesPerPkl without using a loop:

Updated Code with `Buffer.BlockCopy` for `framesPerPkl`:

using System;
using System.IO;
using System.Linq;
using System.Runtime.CompilerServices;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using UnityEngine;

public class KeypointLoader : MonoBehaviour
{
    private NativeArray<Vector3> keypointData;
    private int pklCount;
    private int[] framesPerPkl;

    public async void LoadKeypointData(string path)
    {
        // Dispose existing NativeArray if necessary
        if (keypointData.IsCreated)
            keypointData.Dispose();

        // Read the binary file into a byte array asynchronously
        var data = await File.ReadAllBytesAsync(path);

        // Read .pkl count (unsigned int, 4 bytes)
        pklCount = BitConverter.ToInt32(data, 0);

        // Calculate the starting pointer for framesPerPkl
        int pointer = 4;  // Since we've already read the first 4 bytes for pklCount

        // Allocate the framesPerPkl array
        framesPerPkl = new int[pklCount];

        // Use Buffer.BlockCopy to copy the frame counts directly from the byte array
        Buffer.BlockCopy(data, pointer, framesPerPkl, 0, pklCount * sizeof(int));

        // Move the pointer forward by the amount of data copied (4 bytes per int)
        pointer += pklCount * sizeof(int);

        // Calculate total frame count and data size
        var totalFrames = framesPerPkl.Sum();
        int totalFloats = totalFrames * 51;  // 51 floats per frame

        // Allocate NativeArray to hold all the Vector3 data (floats / 3)
        keypointData = new NativeArray<Vector3>(totalFloats / 3, Allocator.Persistent);

        // Create Memory<byte> to represent the float data slice starting at `pointer`
        var memory = new Memory<byte>(data, pointer, data.Length - pointer);

        unsafe
        {
            // Pin the memory and copy it directly into the NativeArray
            using (var handle = memory.Pin())
            {
                UnsafeUtility.MemCpy(keypointData.GetUnsafePtr(), (void*)handle.Pointer, memory.Length);
            }
        }

        Debug.Log($"Successfully loaded {pklCount} .pkl files with a total of {totalFrames} frames.");
    }

    private void OnDestroy()
    {
        // Dispose of the NativeArray to avoid memory leaks
        if (keypointData.IsCreated)
        {
            keypointData.Dispose();
        }
    }
}

Key Changes:

Buffer.BlockCopy:
- Buffer.BlockCopy(data, pointer, framesPerPkl, 0, pklCount * sizeof(int)) is used to copy the block of memory containing the frame counts from the data byte array directly into the framesPerPkl array.
- This eliminates the need for a loop and copies the memory in one operation, which is more efficient.
Pointer Adjustment:
- After the memory copy, the pointer is advanced by pklCount * sizeof(int) to account for the number of bytes copied (4 bytes per int).

Explanation:

Buffer.BlockCopy is an efficient way to copy blocks of primitive data types (like int or float) from one array to another. In this case, we use it to copy the frame counts (stored as 4-byte ints) from the data byte array into the framesPerPkl integer array in a single call.
This approach avoids the need for a loop, reducing overhead and potentially speeding up the process when dealing with a large number of .pkl files.

Performance Improvement:

Reduced Loop Overhead: By replacing the loop with Buffer.BlockCopy, we reduce the overhead of iterating over the data and avoid multiple calls to BitConverter.ToInt32.
Faster Memory Copy: Buffer.BlockCopy uses native memory copying routines that are highly optimized for transferring large blocks of data, which should be faster than copying one integer at a time.

This modification will further improve the efficiency of your data-loading code by minimizing per-frame processing and leveraging direct memory operations.

Yes, using Sum() in certain scenarios can be inefficient compared to alternatives, particularly in performance-critical code or large-scale data processing. While Sum() is simple and convenient, it introduces overhead because:

It involves iteration: Sum() iterates over the entire collection (in this case, framesPerPkl). For small collections, this overhead is negligible, but for large collections, it can add unnecessary processing time.
Inlining/Function Call Overhead: While modern compilers and runtimes (like .NET) can inline simple methods, there’s still a slight overhead compared to directly using a loop to accumulate values.

Alternatives to `Sum()`

For better performance, especially in performance-critical code, it’s generally more efficient to manually sum values with a loop. This removes the overhead of a method call and allows more optimization.

Example: Manual Summation

Here’s how you can replace framesPerPkl.Sum() with a more efficient manual summation:

// Calculate total frame count manually
int totalFrames = 0;
for (int i = 0; i < framesPerPkl.Length; i++)
{
    totalFrames += framesPerPkl[i];
}

Why is a manual sum potentially more efficient?

No extra method call: By directly summing the values in a loop, you’re removing the need for an additional method call to Sum(), which incurs a small overhead.
Inline loop: The manual summation loop might be more efficiently optimized by the compiler, especially in tight loops or hot paths.

When Does `Sum()` Matter?

Small Data: For small data sets, the performance difference between Sum() and a manual summation will be negligible.
Large Data: For large arrays or performance-critical applications, replacing Sum() with a manual loop can make a noticeable difference because of the lower overhead of a loop compared to a method call and iteration.

Conclusion

Sum() is convenient and generally fine for most applications. However, if performance is critical or you’re dealing with large amounts of data, manually summing the values will typically be more efficient.

If you’re optimizing for performance in a large dataset scenario (like loading keypoints across many .pkl files), replacing Sum() with a loop is a good idea.

To efficiently access the subset of Vector3 elements corresponding to a specific .pkl file within the contiguous NativeArray<Vector3>, you will need to calculate the starting index for the selected .pkl. This can be done by summing the number of frames in the previous .pkl files.

Options for Calculating the Starting Index

1. Prefix Sum (Cumulative Sum)

The prefix sum (or cumulative sum) approach precomputes the starting index for each .pkl file in a separate array. This allows you to directly access the start of any .pkl with a single array lookup, making it the fastest access method once the prefix sum array is built.

2. Looping Over the Frames Count

A loop over the framesPerPkl array to calculate the starting index on the fly is simpler but slightly less efficient than using a prefix sum because it involves recalculating the sum each time you need to access a .pkl.

Which Is Better?

Prefix Sum is the best choice if you need to repeatedly access different .pkl files. It has a slight upfront cost to calculate the starting positions for all .pkl files, but then you get constant-time access (O(1)) to the starting index for any .pkl.
Looping is a fine choice if you access .pkl files infrequently or if you’re concerned about memory overhead and the additional array storage is a concern. The complexity of this method is O(n) where n is the number of .pkl files.

Implementation of Prefix Sum

To implement the prefix sum, you can precompute the starting index of each .pkl and store it in a separate array. When you want to access a specific .pkl, you just look up its starting index from this array.

Step-by-Step Prefix Sum Implementation

Precompute the Prefix Sum: Calculate the starting position of each .pkl in the NativeArray<Vector3>.
Slice the Array: Use the precomputed index to slice into the NativeArray<Vector3>.

Here’s the implementation:

using Unity.Collections;
using UnityEngine;

public class KeypointLoader : MonoBehaviour
{
    private NativeArray<Vector3> keypointData;
    private int[] framesPerPkl;
    private int[] prefixSum;  // Stores the starting index of each .pkl in the keypointData array

    public void InitializePrefixSum()
    {
        // Calculate the prefix sum for efficient access to specific .pkl data
        prefixSum = new int[framesPerPkl.Length];

        int cumulativeSum = 0;
        for (int i = 0; i < framesPerPkl.Length; i++)
        {
            prefixSum[i] = cumulativeSum;
            cumulativeSum += framesPerPkl[i] * 51;  // Each frame has 51 floats (17 joints * 3 coords)
        }
    }

    public NativeSlice<Vector3> GetPklData(int pklIndex)
    {
        // Get the starting index of the selected pkl
        int start = prefixSum[pklIndex] / 3;  // Divide by 3 to convert from float index to Vector3 index
        int frameCount = framesPerPkl[pklIndex];
        int count = frameCount * 17;  // Each frame has 17 Vector3 elements (17 joints)

        // Return a slice of the keypointData array representing the selected .pkl's data
        return new NativeSlice<Vector3>(keypointData, start, count);
    }
}

Explanation:

Prefix Sum Calculation (InitializePrefixSum):
- prefixSum[i] stores the total number of float values (51 floats per frame) that need to be skipped to reach the data for .pkl[i].
- This gives constant-time access to the starting index of any .pkl.
Getting the Slice (GetPklData):
- The starting index for a .pkl is retrieved from prefixSum.
- The number of frames for that .pkl is retrieved from framesPerPkl, and each frame has 17 joints, which means 17 Vector3 elements per frame.
- The result is a NativeSlice<Vector3> that refers to the section of NativeArray<Vector3> corresponding to the specified .pkl.

Why This is Efficient:

Constant Time Access: Once the prefix sum array is built, retrieving the data for any .pkl is done in constant time O(1).
Efficient Slicing: Unity’s NativeSlice<T> allows you to create views into NativeArray<T> without copying the data, so the slice operation is fast and memory-efficient.
No Redundant Calculation: By computing the prefix sum once, you avoid recalculating the starting index every time you access a .pkl.

Conclusion:

Prefix Sum is the most efficient approach for repeatedly accessing different .pkl files, as it allows constant-time access to the starting indices.
The loop-based approach is simpler but less efficient for frequent access because it recalculates the starting index each time.
Using NativeSlice<Vector3> provides a memory-efficient way to view subsets of your data without creating unnecessary allocations or copies.

To create the auxiliary file in Python that identifies each .pkl file with a dance style based on its filename prefix, we will use the table of dance genres and their corresponding prefixes.

Given the mapping you’ve shared, here is the implementation plan:

Read all .pkl filenames.
Match each filename prefix with the dance genre using the provided prefix-to-style mapping.
Write the results to an auxiliary file (e.g., a CSV or a JSON file).

Let’s assume .pkl files have a prefix like gBR_xyz.pkl, where gBR maps to “Break” (as per the table you’ve shared).

Here’s how you can implement it in Python:

Dance Genre Prefix Mapping:

import os
import csv

# Prefix to dance style mapping based on the provided image
dance_genre_map = {
    "gBR": "Break",
    "gPO": "Pop",
    "gLO": "Lock",
    "gMH": "Middle Hip-hop",
    "gLH": "LA style Hip-hop",
    "gHO": "House",
    "gWA": "Waack",
    "gKR": "Krump",
    "gJS": "Street Jazz",
    "gJB": "Ballet Jazz"
}

# Path to the folder containing .pkl files
folder_path = "path_to_your_pkl_files"

def identify_dance_style_for_files(folder_path):
    # List to store the result
    dance_files_info = []

    # Iterate over all .pkl files in the directory
    for filename in os.listdir(folder_path):
        if filename.endswith(".pkl"):
            # Extract the prefix from the filename (assuming first 3 characters)
            prefix = filename[:3]

            # Identify the dance style from the prefix
            dance_style = dance_genre_map.get(prefix, "Unknown")

            # Add to the result list (filename, dance style)
            dance_files_info.append([filename, dance_style])

    # Write the results to a CSV file
    output_file = "dance_styles.csv"
    with open(output_file, mode='w', newline='') as file:
        writer = csv.writer(file)
        writer.writerow(["Filename", "Dance Style"])
        writer.writerows(dance_files_info)

    print(f"Auxiliary file '{output_file}' generated successfully.")

# Run the function
identify_dance_style_for_files(folder_path)

Explanation:

Mapping: The dance_genre_map dictionary maps the filename prefixes (e.g., gBR) to the corresponding dance styles (e.g., Break).
Filename Processing: The script iterates over all .pkl files in the specified folder, extracts the first three characters as the prefix, and looks up the corresponding dance style.
CSV Output: The dance style and filename are written into a CSV file (dance_styles.csv) with two columns: Filename and Dance Style.

Example Output (CSV):

Filename,Dance Style
gBR_001.pkl,Break
gPO_002.pkl,Pop
gLO_003.pkl,Lock
...

This script will help you create an auxiliary file that maps .pkl filenames to their respective dance styles based on the provided prefix-to-genre table.

You can modify the folder path and run the script to generate the auxiliary file. Let me know if you’d prefer a different output format, such as JSON!

To achieve efficient access to specific clips using a combination of enum values (representing dance styles) and index (representing the frame within that style) in Unity, you can implement the following approach:

Steps:

Create an Enum for Dance Styles: This enum will represent all the dance genres based on the prefix (like gBR for Break, gPO for Pop, etc.).
Use a Dictionary to Store the Prefix Sum Indexes: Store the starting position of each genre’s clip in a dictionary where the keys are enums, and values are the start index for each genre in the NativeArray<Vector3>.
Access Clips by Combining Enum and Index: Given an enum value (representing the dance style) and an index (within that dance style), you can compute the global index in the NativeArray<Vector3> and retrieve the corresponding clip.

1. Enum Definition

Define an enum in Unity to represent each dance genre.

public enum DanceStyle
{
    Break,      // gBR
    Pop,        // gPO
    Lock,       // gLO
    MiddleHipHop, // gMH
    LAStyleHipHop, // gLH
    House,      // gHO
    Waack,      // gWA
    Krump,      // gKR
    StreetJazz, // gJS
    BalletJazz  // gJB
}

2. Prefix Sum and Clip Management

Now, calculate the prefix sums to store the starting index of each genre’s data in the NativeArray<Vector3>. This way, combining the dance style enum and frame index will give you the exact location in the array.

using Unity.Collections;
using UnityEngine;
using System.Collections.Generic;

public class KeypointClipLoader : MonoBehaviour
{
    private NativeArray<Vector3> keypointData; // All keypoints data in one contiguous NativeArray
    private Dictionary<DanceStyle, int> prefixSum;  // Stores start index of each dance style
    private Dictionary<DanceStyle, int> styleFrameCounts;  // Stores the number of frames for each style

    // Assume you've initialized keypointData and styleFrameCounts elsewhere
    public void InitializePrefixSum()
    {
        // Initialize the prefix sum dictionary
        prefixSum = new Dictionary<DanceStyle, int>();

        int cumulativeSum = 0;
        foreach (DanceStyle style in System.Enum.GetValues(typeof(DanceStyle)))
        {
            prefixSum[style] = cumulativeSum;
            cumulativeSum += styleFrameCounts[style] * 17;  // Each frame has 17 Vector3 joints
        }
    }

    // Method to access the clip for a given dance style and frame index
    public NativeSlice<Vector3> GetClip(DanceStyle style, int frameIndex)
    {
        if (!prefixSum.ContainsKey(style))
        {
            Debug.LogError("Invalid dance style selected.");
            return default;
        }

        int styleStartIndex = prefixSum[style];  // Get the starting index for the given dance style
        int frameCount = styleFrameCounts[style];

        if (frameIndex < 0 || frameIndex >= frameCount)
        {
            Debug.LogError("Frame index out of bounds.");
            return default;
        }

        // Calculate the starting index for the desired frame
        int startIndex = styleStartIndex + frameIndex * 17;  // 17 Vector3s per frame (17 joints)

        // Return a slice of the NativeArray representing the frame's keypoints
        return new NativeSlice<Vector3>(keypointData, startIndex, 17);
    }

    private void OnDestroy()
    {
        // Dispose of the NativeArray to avoid memory leaks
        if (keypointData.IsCreated)
        {
            keypointData.Dispose();
        }
    }
}

3. Explanation of the Code:

Enum: The DanceStyle enum is used to represent the dance genres.
Prefix Sum Calculation (InitializePrefixSum):
- We compute and store the prefix sum for each dance genre. The prefixSum dictionary maps each DanceStyle to the starting index in the NativeArray<Vector3>.
Clip Access (GetClip):
- By combining the enum value and index, you can calculate the starting position for any frame within the style.
- The NativeSlice<Vector3> is then used to slice into the NativeArray<Vector3> and return the exact data for the frame.
- The returned NativeSlice<Vector3> is memory-efficient, as it doesn’t copy the data, but simply provides a “view” into the NativeArray.

4. Data Access Example

public class DanceClipTester : MonoBehaviour
{
    public KeypointClipLoader clipLoader;

    void Start()
    {
        // Example: Access the 10th frame of the 'Pop' dance style
        NativeSlice<Vector3> clip = clipLoader.GetClip(DanceStyle.Pop, 10);

        // Process the clip data, for example, by drawing lines between joints (Debug.DrawLine example)
        for (int i = 0; i < clip.Length - 1; i++)
        {
            Debug.DrawLine(clip[i], clip[i + 1], Color.green);
        }
    }
}

5. Connecting Joints with `Debug.DrawLine` (COCO Format):

You can use the COCO skeleton structure to connect the joints using Debug.DrawLine. Here’s an example of how you can map the joints and connect them:

// COCO joint connections for a 17-joint skeleton
private int[,] cocoSkeletonConnections = new int[,]
{
    {0, 1}, {0, 2}, {1, 3}, {2, 4},        // Head to Eyes and Ears
    {0, 5}, {0, 6}, {5, 7}, {7, 9},        // Body and left arm
    {6, 8}, {8, 10},                       // Body and right arm
    {5, 11}, {6, 12}, {11, 12},            // Spine and hips
    {11, 13}, {13, 15},                    // Left leg
    {12, 14}, {14, 16}                     // Right leg
};

public void DrawSkeleton(NativeSlice<Vector3> keypoints)
{
    for (int i = 0; i < cocoSkeletonConnections.GetLength(0); i++)
    {
        int joint1 = cocoSkeletonConnections[i, 0];
        int joint2 = cocoSkeletonConnections[i, 1];

        // Draw a line between the two joints
        Debug.DrawLine(keypoints[joint1], keypoints[joint2], Color.green);
    }
}

Summary:

The prefix sum technique allows for fast lookup of any frame within a specific dance style using an enum and index combination.
The NativeSlice<Vector3> ensures memory-efficient access to the specific frames without copying data.
Debug.DrawLine is used to visualize the joints based on the COCO format, where the connections are defined in a 2D array.

This setup will allow you to efficiently access, visualize, and process specific clips from a large dataset of dance keypoints.

Yes, you’re absolutely correct! Since the files are ordered by dance style, it would be much easier and more efficient to store the start and end index (or just the start and count) of each dance style. This way, you can quickly jump to the relevant section of the large array using these indices without needing to calculate the cumulative sum every time you access a new clip.

Benefits of Storing Start and End Indices:

Simpler: You only need to calculate the starting index and the number of frames for each style once.
Faster Lookups: Instead of recalculating or using a cumulative sum, you directly access the precomputed start and end indices to retrieve the slice of the NativeArray.
Memory Efficient: No need to store cumulative sums; you only store two numbers (start and count) per style.

Updated Approach:

Precompute and store the start index and frame count for each dance style.
Use the start and frame count to directly access the correct slice of data.

Code Example

Here’s the updated version of the code using start and end index (or start and count) for each dance style:

1. Storing Start and Count for Each Style

using Unity.Collections;
using UnityEngine;
using System.Collections.Generic;

public class KeypointClipLoader : MonoBehaviour
{
    private NativeArray<Vector3> keypointData; // All keypoints data in one contiguous NativeArray
    private Dictionary<DanceStyle, (int startIndex, int frameCount)> danceStyleIndexMap;

    // Assume you've initialized keypointData and filled the dance style index information elsewhere
    public void InitializeDanceStyleIndexMap(Dictionary<DanceStyle, int> styleFrameCounts)
    {
        danceStyleIndexMap = new Dictionary<DanceStyle, (int startIndex, int frameCount)>();

        int cumulativeSum = 0;
        foreach (DanceStyle style in System.Enum.GetValues(typeof(DanceStyle)))
        {
            int frameCount = styleFrameCounts[style];
            danceStyleIndexMap[style] = (cumulativeSum, frameCount);
            cumulativeSum += frameCount * 17;  // Each frame has 17 Vector3 joints
        }
    }

    // Method to access the clip for a given dance style and frame index
    public NativeSlice<Vector3> GetClip(DanceStyle style, int frameIndex)
    {
        if (!danceStyleIndexMap.ContainsKey(style))
        {
            Debug.LogError("Invalid dance style selected.");
            return default;
        }

        // Get the start index and frame count for the selected dance style
        var (startIndex, frameCount) = danceStyleIndexMap[style];

        if (frameIndex < 0 || frameIndex >= frameCount)
        {
            Debug.LogError("Frame index out of bounds.");
            return default;
        }

        // Calculate the starting index for the desired frame
        int sliceStartIndex = startIndex + frameIndex * 17;  // 17 Vector3s per frame (17 joints)

        // Return a slice of the NativeArray representing the selected frame's keypoints
        return new NativeSlice<Vector3>(keypointData, sliceStartIndex, 17);
    }

    private void OnDestroy()
    {
        // Dispose of the NativeArray to avoid memory leaks
        if (keypointData.IsCreated)
        {
            keypointData.Dispose();
        }
    }
}

2. Explanation:

Dance Style Index Map:
- We store a dictionary (danceStyleIndexMap) where each DanceStyle maps to a tuple containing the startIndex and frameCount for that style.
Initialize Map:
- When you initialize the loader, you calculate the starting index for each dance style based on the cumulative number of frames and store it in the dictionary.
- You also store the number of frames (frameCount) for each style.
Fast Access:
- When accessing a clip, you retrieve the startIndex and frameCount for the requested dance style from the dictionary, and use this to calculate the correct slice of the NativeArray<Vector3>.
Efficient:
- This avoids recalculating cumulative sums every time you access a frame and provides constant-time access to any frame within the style.

3. Usage Example:

public class DanceClipTester : MonoBehaviour
{
    public KeypointClipLoader clipLoader;

    void Start()
    {
        // Example: Access the 10th frame of the 'Pop' dance style
        NativeSlice<Vector3> clip = clipLoader.GetClip(DanceStyle.Pop, 10);

        // Process the clip data, for example, by drawing lines between joints
        clipLoader.DrawSkeleton(clip);
    }
}

4. Connecting Joints with `Debug.DrawLine`:

To visualize the skeleton using Debug.DrawLine (based on the COCO format), you can create connections between joints as shown in the earlier code:

private int[,] cocoSkeletonConnections = new int[,]
{
    {0, 1}, {0, 2}, {1, 3}, {2, 4},        // Head to Eyes and Ears
    {0, 5}, {0, 6}, {5, 7}, {7, 9},        // Body and left arm
    {6, 8}, {8, 10},                       // Body and right arm
    {5, 11}, {6, 12}, {11, 12},            // Spine and hips
    {11, 13}, {13, 15},                    // Left leg
    {12, 14}, {14, 16}                     // Right leg
};

public void DrawSkeleton(NativeSlice<Vector3> keypoints)
{
    for (int i = 0; i < cocoSkeletonConnections.GetLength(0); i++)
    {
        int joint1 = cocoSkeletonConnections[i, 0];
        int joint2 = cocoSkeletonConnections[i, 1];

        // Draw a line between the two joints
        Debug.DrawLine(keypoints[joint1], keypoints[joint2], Color.green);
    }
}

Summary:

Simple and Efficient: Storing just the start and frame count for each dance style simplifies access and allows constant-time lookups without needing to compute cumulative sums.
Direct Memory Access: Using NativeSlice<Vector3> ensures efficient memory access without copying data.
Visualization: Using Debug.DrawLine to connect the joints allows for easy visual debugging and rendering of the skeleton based on the COCO format.

This approach provides a clean and efficient solution for managing and visualizing large keypoint datasets in Unity, especially when the .pkl files are ordered by dance style.

initialize the python dict with the dance names and counter of 0

def identify_dance_style_for_files(folder_path): # List to store the result dance_files_info = {}

# Iterate over all .pkl files in the directory
for filename in os.listdir(folder_path):
    if filename.endswith(".pkl"):

        # Extract the prefix from the filename (assuming first 3 characters)
        prefix = filename[:3]

        # Identify the dance style from the prefix
        dance_style = dance_genre_map.get(prefix, "Unknown")

        # Add to the result list (filename, dance style)
        dance_files_info[dance_style] += 1

# Write the results to a CSV file
output_file = "dance_styles.json"

with open(output_file, mode='w', newline='') as file:
    json.dump(dance_files_info)

print(f"Auxiliary file '{output_file}' generated successfully.")

To initialize the dance_files_info dictionary with the dance names as keys and a counter of 0 for each, you can modify your code like this:

Full Code with the Initialization

import os
import json

# Prefix to dance style mapping based on the provided image
dance_genre_map = {
    "gBR": "Break",
    "gPO": "Pop",
    "gLO": "Lock",
    "gMH": "Middle Hip-hop",
    "gLH": "LA style Hip-hop",
    "gHO": "House",
    "gWA": "Waack",
    "gKR": "Krump",
    "gJS": "Street Jazz",
    "gJB": "Ballet Jazz"
}

# Initialize the dictionary with the dance styles and a counter of 0
dance_files_info = {style: 0 for style in dance_genre_map.values()}

def identify_dance_style_for_files(folder_path):
    # Iterate over all .pkl files in the directory
    for filename in os.listdir(folder_path):
        if filename.endswith(".pkl"):

            # Extract the prefix from the filename (assuming first 3 characters)
            prefix = filename[:3]

            # Identify the dance style from the prefix
            dance_style = dance_genre_map.get(prefix, "Unknown")

            # Increment the counter for the corresponding dance style
            if dance_style != "Unknown":
                dance_files_info[dance_style] += 1

    # Write the results to a JSON file
    output_file = "dance_styles.json"

    with open(output_file, mode='w', newline='') as file:
        json.dump(dance_files_info, file, indent=4)

    print(f"Auxiliary file '{output_file}' generated successfully.")

# Example of how to call the function
folder_path = "path_to_your_pkl_files"
identify_dance_style_for_files(folder_path)

Key Changes:

Initialize dance_files_info with 0:
- The dance_files_info dictionary is initialized using a dictionary comprehension that maps each dance genre (from dance_genre_map.values()) to an initial count of 0.
Check for Unknown Styles:
- If a file’s prefix does not map to a known dance style, it will not update the dictionary (you could add handling for unknowns if needed).
Write to JSON:
- After processing all .pkl files, the result (dance_files_info) is saved into a JSON file (dance_styles.json), with pretty printing (indent=4) for better readability.

Example JSON Output:

{
    "Break": 5,
    "Pop": 3,
    "Lock": 2,
    "Middle Hip-hop": 1,
    "LA style Hip-hop": 4,
    "House": 0,
    "Waack": 0,
    "Krump": 0,
    "Street Jazz": 0,
    "Ballet Jazz": 0
}

This code initializes the counter for each dance style with 0, processes the .pkl files to count how many belong to each style, and then writes the results to a JSON file.

To parse the dance_styles.json file in Unity, build a second prefix sum based on the counts of each genre, and use that to efficiently index into your NativeArray<Vector3>, follow these steps.

Step-by-Step Implementation

Parse dance_styles.json: We’ll load the JSON file in Unity and deserialize it into a dictionary that maps dance styles to their respective counts.
Build a Prefix Sum: We compute a prefix sum array that gives us the starting index for each dance style within the NativeArray<Vector3>.
Access Data Using the Prefix Sum: Using the prefix sum, you can quickly access the starting position of a dance style in the NativeArray.

Step 1: Parsing the JSON File in Unity

First, let’s write the Unity code to load and parse the dance_styles.json file.

`dance_styles.json` Example:

{
    "Break": 5,
    "Pop": 3,
    "Lock": 2,
    "Middle Hip-hop": 1,
    "LA style Hip-hop": 4,
    "House": 0,
    "Waack": 0,
    "Krump": 0,
    "Street Jazz": 0,
    "Ballet Jazz": 0
}

Unity Script to Parse the JSON

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

public class DanceStyleParser : MonoBehaviour
{
    private Dictionary<string, int> danceStyleCounts;

    // Path to the JSON file (assuming it's in the Resources folder)
    private string jsonFilePath = "Assets/dance_styles.json";

    public void ParseDanceStyles()
    {
        if (!File.Exists(jsonFilePath))
        {
            Debug.LogError("JSON file not found at path: " + jsonFilePath);
            return;
        }

        // Read the JSON file
        string jsonData = File.ReadAllText(jsonFilePath);

        // Deserialize the JSON data into a dictionary
        danceStyleCounts = JsonUtility.FromJson<DictionaryWrapper>(jsonData).ToDictionary();
    }

    // A helper class to hold the dictionary since Unity's JsonUtility needs a wrapper
    [System.Serializable]
    public class DictionaryWrapper
    {
        public List<string> keys = new List<string>();
        public List<int> values = new List<int>();

        public Dictionary<string, int> ToDictionary()
        {
            var dict = new Dictionary<string, int>();
            for (int i = 0; i < keys.Count; i++)
            {
                dict.Add(keys[i], values[i]);
            }
            return dict;
        }
    }
}

Step 2: Build a Second Prefix Sum Based on Genres

Once we have the parsed data, we can create a prefix sum of genres. This will give us the starting index of each dance style in the NativeArray<Vector3>.

Prefix Sum Construction

using System.Collections.Generic;
using Unity.Collections;
using UnityEngine;

public class KeypointClipLoader : MonoBehaviour
{
    private NativeArray<Vector3> keypointData;  // Assume this is filled elsewhere
    private Dictionary<string, (int startIndex, int count)> genrePrefixSum;

    public void BuildPrefixSum(Dictionary<string, int> danceStyleCounts)
    {
        genrePrefixSum = new Dictionary<string, (int startIndex, int count)>();

        int currentIndex = 0;
        foreach (var genre in danceStyleCounts)
        {
            int frameCount = genre.Value;
            genrePrefixSum[genre.Key] = (currentIndex, frameCount);
            currentIndex += frameCount * 17;  // Each frame has 17 Vector3 joints
        }
    }

    // Method to get a specific clip based on the genre and frame index
    public NativeSlice<Vector3> GetClip(string genre, int frameIndex)
    {
        if (!genrePrefixSum.ContainsKey(genre))
        {
            Debug.LogError("Invalid genre selected: " + genre);
            return default;
        }

        var (startIndex, frameCount) = genrePrefixSum[genre];

        if (frameIndex < 0 || frameIndex >= frameCount)
        {
            Debug.LogError("Frame index out of bounds for genre: " + genre);
            return default;
        }

        // Calculate the starting index for the frame
        int sliceStartIndex = startIndex + frameIndex * 17;  // 17 Vector3s per frame (17 joints)

        // Return a slice of the NativeArray representing the selected frame's keypoints
        return new NativeSlice<Vector3>(keypointData, sliceStartIndex, 17);
    }

    private void OnDestroy()
    {
        // Dispose of the NativeArray to avoid memory leaks
        if (keypointData.IsCreated)
        {
            keypointData.Dispose();
        }
    }
}

Step 3: Usage Example

Now that we have the prefix sum, we can easily access data for any genre and frame:

public class DanceClipTester : MonoBehaviour
{
    public KeypointClipLoader clipLoader;
    public DanceStyleParser parser;

    void Start()
    {
        // Parse the JSON file
        parser.ParseDanceStyles();

        // Build the prefix sum based on the parsed dance styles
        clipLoader.BuildPrefixSum(parser.danceStyleCounts);

        // Example: Access the 2nd frame of the 'Pop' dance style
        NativeSlice<Vector3> clip = clipLoader.GetClip("Pop", 1);

        // Process the clip data, for example, by drawing lines between joints
        clipLoader.DrawSkeleton(clip);
    }
}

4. Drawing Skeletons Based on Keypoints

Here’s how to visualize the skeleton using Debug.DrawLine based on the COCO format:

private int[,] cocoSkeletonConnections = new int[,]
{
    {0, 1}, {0, 2}, {1, 3}, {2, 4},        // Head to Eyes and Ears
    {0, 5}, {0, 6}, {5, 7}, {7, 9},        // Body and left arm
    {6, 8}, {8, 10},                       // Body and right arm
    {5, 11}, {6, 12}, {11, 12},            // Spine and hips
    {11, 13}, {13, 15},                    // Left leg
    {12, 14}, {14, 16}                     // Right leg
};

public void DrawSkeleton(NativeSlice<Vector3> keypoints)
{
    for (int i = 0; i < cocoSkeletonConnections.GetLength(0); i++)
    {
        int joint1 = cocoSkeletonConnections[i, 0];
        int joint2 = cocoSkeletonConnections[i, 1];

        // Draw a line between the two joints
        Debug.DrawLine(keypoints[joint1], keypoints[joint2], Color.green);
    }
}

Summary:

JSON Parsing: The dance_styles.json file is parsed into a dictionary that maps each dance genre to its frame count.
Prefix Sum: A second prefix sum is created that gives the start and count of frames for each genre.
Efficient Access: Using this prefix sum, you can quickly retrieve the frames for any dance genre and index into the large NativeArray<Vector3>.
Visualization: Using Debug.DrawLine with the COCO skeleton connections allows you to visualize the keypoints as a skeleton in Unity.

This approach efficiently handles large amounts of data and allows for easy access to specific dance genres and frames in Unity.

The issue you’re encountering is related to how Unity’s JsonUtility works. Unity’s JsonUtility does not handle generic dictionaries directly or parse JSON objects where the keys are dynamic (like the example you provided). This is a known limitation of JsonUtility.

To solve this, you’ll need to either:

Manually parse the JSON using another method, such as the Newtonsoft.Json (JSON.NET) library, which handles dictionaries.
Work around Unity’s limitations by creating a wrapper class that can be parsed by JsonUtility and then converted into a dictionary.

Option 1: Use `Newtonsoft.Json` (JSON.NET)

Newtonsoft.Json (JSON.NET) is a popular library for handling JSON in Unity, and it can easily handle dictionaries and more complex JSON structures.

Steps to Install Newtonsoft.Json in Unity:

Install the Newtonsoft.Json package from Unity’s package manager:
- Open Window > Package Manager.
- Search for “Json.NET” or “Newtonsoft.Json” and install it.
Modify Your Code to Use Newtonsoft.Json:

using System.Collections.Generic;
using Newtonsoft.Json;  // Add the Newtonsoft.Json namespace
using System.IO;
using UnityEngine;

public class DanceStyleParser : MonoBehaviour
{
    private Dictionary<string, int> danceStyleCounts;

    // Path to the JSON file
    private string jsonFilePath = "Assets/dance_styles.json";

    public void ParseDanceStyles()
    {
        if (!File.Exists(jsonFilePath))
        {
            Debug.LogError("JSON file not found at path: " + jsonFilePath);
            return;
        }

        // Read the JSON file
        string jsonData = File.ReadAllText(jsonFilePath);

        // Deserialize the JSON data into a dictionary using Newtonsoft.Json
        danceStyleCounts = JsonConvert.DeserializeObject<Dictionary<string, int>>(jsonData);

        foreach (var entry in danceStyleCounts)
        {
            Debug.Log($"Dance Style: {entry.Key}, Count: {entry.Value}");
        }
    }
}

Explanation:

Newtonsoft.Json is used to deserialize the JSON data into a Dictionary<string, int> directly.
This approach handles JSON objects with dynamic keys (like the example you provided) without the limitations of JsonUtility.

Option 2: Work Around `JsonUtility` Limitation

If you prefer not to use third-party libraries, you can create a workaround by manually converting your JSON into a format that JsonUtility can handle. This involves creating a wrapper class to represent the data.

Wrapper Class for Unity’s `JsonUtility`

Modify JSON Format: You would need to modify the JSON structure to fit a format that JsonUtility can parse. Instead of a dictionary, represent it as an array of key-value pairs.
Modified JSON Example:

{
    "danceStyles": [
        {"style": "Break", "count": 141},
        {"style": "Pop", "count": 140},
        {"style": "Lock", "count": 141},
        {"style": "Middle Hip-hop", "count": 141},
        {"style": "LA style Hip-hop", "count": 141},
        {"style": "House", "count": 141},
        {"style": "Waack", "count": 140},
        {"style": "Krump", "count": 141},
        {"style": "Street Jazz", "count": 141},
        {"style": "Ballet Jazz", "count": 141}
    ]
}

Create Wrapper Classes:

using System.Collections.Generic;
using UnityEngine;

public class DanceStyleParser : MonoBehaviour
{
    private Dictionary<string, int> danceStyleCounts = new Dictionary<string, int>();

    // Path to the JSON file (assuming it's in Resources folder or somewhere else)
    private string jsonFilePath = "Assets/dance_styles.json";

    public void ParseDanceStyles()
    {
        if (!System.IO.File.Exists(jsonFilePath))
        {
            Debug.LogError("JSON file not found at path: " + jsonFilePath);
            return;
        }

        // Read the JSON file
        string jsonData = System.IO.File.ReadAllText(jsonFilePath);

        // Deserialize the JSON data into a wrapper class
        DanceStyleListWrapper wrapper = JsonUtility.FromJson<DanceStyleListWrapper>(jsonData);

        // Convert the wrapper data into a dictionary
        foreach (var entry in wrapper.danceStyles)
        {
            danceStyleCounts.Add(entry.style, entry.count);
            Debug.Log($"Dance Style: {entry.style}, Count: {entry.count}");
        }
    }
}

[System.Serializable]
public class DanceStyleEntry
{
    public string style;
    public int count;
}

[System.Serializable]
public class DanceStyleListWrapper
{
    public List<DanceStyleEntry> danceStyles;
}

Explanation:

DanceStyleEntry: Represents each entry in the JSON (the dance style and its count).
DanceStyleListWrapper: Wraps the list of DanceStyleEntry objects, which JsonUtility can now parse.
danceStyleCounts Dictionary: After parsing, the data is converted into a dictionary for easy access.

Summary:

Option 1 (Newtonsoft.Json): This is the most straightforward and efficient way if you’re okay with using an external library. It allows you to directly deserialize the JSON into a dictionary.
Option 2 (JsonUtility): If you prefer sticking to Unity’s built-in tools, you can use a wrapper class and convert the JSON into a format that JsonUtility can handle. However, this requires modifying the structure of your JSON.

For most Unity projects, using Newtonsoft.Json (Option 1) will save time and be more flexible for handling various JSON structures. If you’re working with Unity’s JsonUtility, you’ll need to use the workaround described in Option 2.

I have a .pkl file generated by this code (WHAM). I want to parse it and visualize it frame-by-frame. i’m interested in pose_world and trans_World

    # ========= Store results ========= #
    pred_body_pose = matrix_to_axis_angle(pred['poses_body']).cpu().numpy().reshape(-1, 69)
    pred_root = matrix_to_axis_angle(pred['poses_root_cam']).cpu().numpy().reshape(-1, 3)
    pred_root_world = matrix_to_axis_angle(pred['poses_root_world']).cpu().numpy().reshape(-1, 3)
    pred_pose = np.concatenate((pred_root, pred_body_pose), axis=-1)
    pred_pose_world = np.concatenate((pred_root_world, pred_body_pose), axis=-1)
    pred_trans = (pred['trans_cam'] - network.output.offset).cpu().numpy()

    results[_id]['pose'] = pred_pose
    results[_id]['trans'] = pred_trans
    results[_id]['pose_world'] = pred_pose_world
    results[_id]['trans_world'] = pred['trans_world'].cpu().squeeze(0).numpy()
    results[_id]['betas'] = pred['betas'].cpu().squeeze(0).numpy()
    results[_id]['verts'] = (pred['verts_cam'] + pred['trans_cam'].unsqueeze(1)).cpu().numpy()
    results[_id]['frame_ids'] = frame_id

if save_pkl:
    joblib.dump(results, osp.join(output_pth, "wham_output.pkl"))

To parse and visualize the pose_world and trans_world data from the .pkl file generated by WHAM, you can follow these steps:

Steps:

Load the .pkl file: The .pkl file is stored in Python using the joblib.dump() method, so you can load it using joblib.load() or pickle.load().
Extract the pose_world and trans_world: For each frame, you’ll extract the pose_world and trans_world data, which are numpy arrays.
Visualize: Use Python libraries such as matplotlib or mpl_toolkits.mplot3d to visualize the poses and translations frame-by-frame.
Display frame-by-frame: You can create an animated visualization that shows the pose_world and trans_world data for each frame.

Step 1: Load the `.pkl` File

You will use joblib to load the .pkl file and access the stored data.

import joblib
import numpy as np

# Load the .pkl file
pkl_path = "path_to_your_pkl_file.pkl"
results = joblib.load(pkl_path)

# Let's check what keys are available for each entry in the results
first_entry = list(results.keys())[0]
print("Available data fields:", results[first_entry].keys())

Step 2: Extract `pose_world` and `trans_world`

The results dictionary contains multiple frames for each result, and each result has the keys like pose_world and trans_world. You can access these for each frame.

# Extract pose_world and trans_world from the results
pose_world_data = []
trans_world_data = []

for frame_id, frame_data in results.items():
    pose_world = frame_data['pose_world']
    trans_world = frame_data['trans_world']

    # Store the data for visualization
    pose_world_data.append(pose_world)
    trans_world_data.append(trans_world)

print(f"Extracted {len(pose_world_data)} frames of pose_world and trans_world data.")

Step 3: Visualize Frame-by-Frame

You can visualize the data frame-by-frame using matplotlib with 3D plotting. Here is a simple way to visualize the pose_world and trans_world for each frame.

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def visualize_frame(pose_world, trans_world):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # Visualize the root translation
    ax.scatter(trans_world[0], trans_world[1], trans_world[2], c='r', marker='o', label="Root Translation")

    # Visualize the body pose (pose_world is a concatenation of root and body, as per your original code)
    # Assuming 3D rotation matrices are applied to the joints:
    for i in range(0, len(pose_world), 3):
        ax.quiver(trans_world[0], trans_world[1], trans_world[2], pose_world[i], pose_world[i+1], pose_world[i+2])

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    plt.legend()
    plt.show()

# Example: visualize the first frame
visualize_frame(pose_world_data[0], trans_world_data[0])

Step 4: Animate Frame-by-Frame Visualization

To animate the poses over time, you can use matplotlib.animation. This will give you a frame-by-frame display of the poses in 3D space.

import matplotlib.animation as animation

def update_frame(num, pose_world_data, trans_world_data, ax):
    ax.clear()

    # Plot the current frame's pose_world and trans_world data
    trans_world = trans_world_data[num]
    pose_world = pose_world_data[num]

    ax.scatter(trans_world[0], trans_world[1], trans_world[2], c='r', marker='o', label="Root Translation")

    for i in range(0, len(pose_world), 3):
        ax.quiver(trans_world[0], trans_world[1], trans_world[2], pose_world[i], pose_world[i+1], pose_world[i+2])

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.legend()

def animate_pose_world(pose_world_data, trans_world_data):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    ani = animation.FuncAnimation(
        fig, update_frame, frames=len(pose_world_data), fargs=(pose_world_data, trans_world_data, ax),
        interval=100  # Adjust the interval for speed
    )

    plt.show()

# Animate the entire sequence
animate_pose_world(pose_world_data, trans_world_data)

Additional Considerations

Pose Representation: If pose_world is in axis-angle format (as your original code suggests), you may need to convert these rotations into a format that can be plotted, such as quaternions or rotation matrices.
Scaling and Adjustment: Depending on the data, you may need to adjust the visualization scaling (e.g., scaling the quivers or adjusting axis limits).
Animation Speed: You can adjust the interval in FuncAnimation to control the speed of the frame-by-frame playback.

Summary:

Load the .pkl file using joblib.
Extract the pose_world and trans_world data from the loaded dictionary.
Visualize each frame using matplotlib with 3D plotting, showing both the pose and root translation.
Animate the poses over time using matplotlib.animation for frame-by-frame playback.

This will allow you to visualize the pose and translation of the body in 3D space based on the data stored in the .pkl file.

Thank you for the clarification! Let’s revise the visualization code based on the correct way of parsing the pose_world and trans_world data. Specifically:

pose_world: This seems to be a list where each entry represents the pose for a single frame.
trans_world: Similarly, this represents the translation for each frame.

Now we will visualize the pose_world and trans_world data frame-by-frame.

Step 1: Load the `.pkl` Data Correctly

Here’s how you will load the .pkl file and iterate over the frames in pose_world and trans_world.

import joblib

# Load the .pkl file
pkl_path = "path_to_your_pkl_file.pkl"
results = joblib.load(pkl_path)

# Extract the first entry's pose_world and trans_world
pose_world = results[0]['pose_world']
trans_world = results[0]['trans_world']

# Check the number of frames
num_frames = len(pose_world)
print(f"Number of frames: {num_frames}")

Step 2: Frame-by-Frame Visualization

We will now iterate through each frame in pose_world and trans_world, and visualize it using matplotlib. We will plot the root translation (trans_world) and body pose (pose_world).

Given the structure, it appears that:

trans_world[i] corresponds to the root translation of the body in the ith frame.
pose_world[i] contains the body pose information for the ith frame.

Updated Visualization Function

Here is the revised visualization function to show the pose and translation for each frame:

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

def visualize_frame(pose, trans):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # Plot the root translation (as a point)
    ax.scatter(trans[0], trans[1], trans[2], c='r', marker='o', label="Root Translation")

    # Plot the body pose (assuming we have 3D vectors for the pose)
    # Assuming pose contains a series of (x, y, z) triplets for joints or directions
    num_joints = len(pose) // 3  # Each joint has 3 values (x, y, z)

    for j in range(num_joints):
        x, y, z = pose[3*j], pose[3*j+1], pose[3*j+2]
        ax.scatter(x + trans[0], y + trans[1], z + trans[2], c='b', marker='^', label=f"Joint {j}" if j == 0 else "")

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    plt.legend()
    plt.show()

# Example: visualize the first frame
visualize_frame(pose_world[0], trans_world[0])

Step 3: Frame-by-Frame Animation

To create an animated visualization that shows each frame in sequence, we can use matplotlib.animation:

import matplotlib.animation as animation

def update_frame(num, pose_world_data, trans_world_data, ax):
    ax.clear()

    # Get the current frame's pose and translation
    pose = pose_world_data[num]
    trans = trans_world_data[num]

    # Plot the root translation (as a point)
    ax.scatter(trans[0], trans[1], trans[2], c='r', marker='o', label="Root Translation")

    # Plot the body pose
    num_joints = len(pose) // 3
    for j in range(num_joints):
        x, y, z = pose[3*j], pose[3*j+1], pose[3*j+2]
        ax.scatter(x + trans[0], y + trans[1], z + trans[2], c='b', marker='^', label=f"Joint {j}" if j == 0 else "")

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.legend()

def animate_pose_world(pose_world_data, trans_world_data):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    ani = animation.FuncAnimation(
        fig, update_frame, frames=len(pose_world_data), fargs=(pose_world_data, trans_world_data, ax),
        interval=100  # Adjust the interval for speed
    )

    plt.show()

# Animate the entire sequence
animate_pose_world(pose_world, trans_world)

Explanation:

Single Frame Visualization:
- The visualize_frame() function takes the pose (pose_world[i]) and translation (trans_world[i]) for a single frame and visualizes it in 3D using matplotlib.
- Each joint in the pose is plotted relative to the root translation (trans_world[i]).
Animated Visualization:
- The animate_pose_world() function visualizes the data frame-by-frame, creating an animation.
- The FuncAnimation() function from matplotlib.animation updates the plot for each frame in the sequence, allowing you to watch the body pose evolve over time.

Key Assumptions:

Pose Data: The pose_world[i] is assumed to be a list of (x, y, z) coordinates for each joint in the body.
Translation: The trans_world[i] is the root translation for the body, applied to all the joints to position them correctly in space.

Notes:

If pose_world is stored in a different format (e.g., rotations or quaternions), you might need to convert them to joint positions before visualization.
You can adjust the speed of the animation by changing the interval in FuncAnimation().

This revised visualization code should now properly display the pose_world and trans_world data frame-by-frame, based on the correct parsing of the .pkl file.

The fact that each pose_world frame has 72 floats (which corresponds to 24 joints with 3D coordinates) suggests that the pose model being used is likely based on a 24-joint human body model. Since you’re using ViTPose (Vision Transformer Pose), we can deduce that the 24 joints represent a standard human skeletal model commonly used in pose estimation.

Given this, the 72 floats can be interpreted as 24 joints × 3 coordinates (x, y, z) per joint. This gives us 72 values per frame, where each joint has 3 floats representing its position in 3D space.

Likely Pose Model

There are several standard human body models used in pose estimation that have 24 joints, and from your description, it’s highly likely that the ViTPose model is based on one of the following models:

SMPL (Skinned Multi-Person Linear Model):
- SMPL is a widely used 3D human body model, which consists of 24 joints. It’s used for generating human poses in 3D space.
- The SMPL model is often used in conjunction with 3D pose estimation networks, and it makes sense that ViTPose could be using it since SMPL fits well with body shape and pose estimation tasks.
H36M (Human3.6M):
- Another common dataset and model, which also uses 24 keypoints for human pose estimation.
- This model is used extensively for training and evaluating pose estimation models, and it is possible that ViTPose could have been trained or evaluated on this dataset.

Since your network produces 72 floats per frame, and ViTPose uses a transformer-based approach, it’s likely adhering to one of these models.

Breakdown of the 24 Joints:

These joints typically cover:

Head and neck: Neck, head
Torso: Spine, hips, shoulders
Arms: Elbows, wrists
Legs: Knees, ankles

Here’s a typical breakdown of 24 joints for SMPL:

Pelvis
Left Hip
Right Hip
Spine 1
Left Knee
Right Knee
Spine 2
Left Ankle
Right Ankle
Spine 3
Left Toe
Right Toe
Neck
Left Shoulder
Right Shoulder
Head
Left Elbow
Right Elbow
Left Wrist
Right Wrist
Left Hand
Right Hand
Nose
Eyes (Optional for some models)

Visualization

Based on this understanding, we can now interpret the 72 floats as corresponding to 24 joints, each with 3D coordinates.

Updated Visualization Code for 24 Joints

Let’s revise the code to reflect the fact that you’re working with 24 joints (72 floats):

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

# Updated function to visualize a single frame of pose and translation
def visualize_frame(pose, trans):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # Plot the root translation (as a point)
    ax.scatter(trans[0], trans[1], trans[2], c='r', marker='o', label="Root Translation")

    # Visualize the body pose (24 joints, 72 floats)
    num_joints = 24  # 24 joints
    for j in range(num_joints):
        x, y, z = pose[3 * j], pose[3 * j + 1], pose[3 * j + 2]
        ax.scatter(x + trans[0], y + trans[1], z + trans[2], c='b', marker='^', label=f"Joint {j}" if j == 0 else "")

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    plt.legend()
    plt.show()

# Example: visualize the first frame
visualize_frame(pose_world[0], trans_world[0])

Summary:

72 floats per frame indicates 24 joints, which corresponds to a typical human body model such as SMPL (Skinned Multi-Person Linear Model) or possibly the H36M dataset.
ViTPose likely uses this 24-joint human body model, with each joint represented by 3 floats (x, y, z) for 3D coordinates.
The revised visualization code now correctly interprets the data as 24 joints, allowing you to visualize the pose frame-by-frame.

This interpretation should align with your model’s output and allow for meaningful visualization of human poses.

The SMPL model defines 24 joints and connects them to form a skeleton through “bones” or connections between joints. To visualize this in 3D, we need to:

Define the connections (bones) between the 24 joints of the SMPL model.
Draw lines between connected joints in the 3D plot to represent bones.

Step-by-Step Approach

Define the Joint Connections (SMPL Skeleton): These connections specify which joints are connected to form the body skeleton.
Update the Visualization: For each frame, we will draw the bones (lines between joints) rather than just displaying the joints.

SMPL Skeleton Joint Connections

Here’s the connection map for the SMPL skeleton model. Each pair represents two connected joints (forming a bone).

0: Pelvis
1: Left Hip
2: Right Hip
3: Spine 1
4: Left Knee
5: Right Knee
6: Spine 2
7: Left Ankle
8: Right Ankle
9: Spine 3
10: Left Toe
11: Right Toe
12: Neck
13: Left Shoulder
14: Right Shoulder
15: Head
16: Left Elbow
17: Right Elbow
18: Left Wrist
19: Right Wrist
20: Left Hand
21: Right Hand
22: Nose
23: Eyes (or an additional joint in some models)

Joint Connections (Bone Structure)

Here’s a common set of connections used for SMPL:

smpl_skeleton = [
    (0, 1), (0, 2), (0, 3),          # Pelvis to hips and spine
    (1, 4), (2, 5),                  # Hips to knees
    (4, 7), (5, 8),                  # Knees to ankles
    (7, 10), (8, 11),                # Ankles to toes
    (3, 6), (6, 9), (9, 12),         # Spine to upper spine to neck
    (12, 13), (12, 14),              # Neck to shoulders
    (13, 16), (14, 17),              # Shoulders to elbows
    (16, 18), (17, 19),              # Elbows to wrists
    (18, 20), (19, 21),              # Wrists to hands
    (9, 15), (15, 22), (22, 23)      # Spine to head, head to nose, nose to eyes
]

Visualization Code

We will now modify the visualization code to draw lines (bones) between connected joints.

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

# Define the SMPL skeleton connections (bones)
smpl_skeleton = [
    (0, 1), (0, 2), (0, 3),          # Pelvis to hips and spine
    (1, 4), (2, 5),                  # Hips to knees
    (4, 7), (5, 8),                  # Knees to ankles
    (7, 10), (8, 11),                # Ankles to toes
    (3, 6), (6, 9), (9, 12),         # Spine to upper spine to neck
    (12, 13), (12, 14),              # Neck to shoulders
    (13, 16), (14, 17),              # Shoulders to elbows
    (16, 18), (17, 19),              # Elbows to wrists
    (18, 20), (19, 21),              # Wrists to hands
    (9, 15), (15, 22), (22, 23)      # Spine to head, head to nose, nose to eyes
]

# Visualize a single frame with the skeleton drawn as bones
def visualize_skeleton_frame(pose, trans):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # Plot the root translation (as a point)
    ax.scatter(trans[0], trans[1], trans[2], c='r', marker='o', label="Root Translation")

    # Plot each bone in the skeleton
    for joint1, joint2 in smpl_skeleton:
        # Get the 3D coordinates of the two joints
        x1, y1, z1 = pose[3 * joint1], pose[3 * joint1 + 1], pose[3 * joint1 + 2]
        x2, y2, z2 = pose[3 * joint2], pose[3 * joint2 + 1], pose[3 * joint2 + 2]

        # Draw the line between the two joints (bone)
        ax.plot(
            [x1 + trans[0], x2 + trans[0]],
            [y1 + trans[1], y2 + trans[1]],
            [z1 + trans[2], z2 + trans[2]],
            color='blue'
        )

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    plt.legend()
    plt.show()

# Example: visualize the first frame
visualize_skeleton_frame(pose_world[0], trans_world[0])

Explanation:

SMPL Skeleton:
- The smpl_skeleton list defines the connections (bones) between joints. Each tuple (joint1, joint2) indicates that there is a bone between joint1 and joint2.
Drawing Bones:
- For each bone, the coordinates of the two connected joints are retrieved from pose (which contains 3D positions of joints), and then a line is drawn between the two joints using ax.plot().
Translation:
- The root translation (trans) is applied to each joint’s position to visualize the pose in world coordinates.

Example:

If pose_world and trans_world represent the full body pose and translation for each frame, this code will visualize the skeleton by drawing bones between joints frame-by-frame.

Animation Example:

You can also create an animation of the skeleton over time, just like before, but now with bones drawn between the joints.

import matplotlib.animation as animation

def update_skeleton_frame(num, pose_world_data, trans_world_data, ax):
    ax.clear()

    pose = pose_world_data[num]
    trans = trans_world_data[num]

    # Plot each bone in the skeleton
    for joint1, joint2 in smpl_skeleton:
        x1, y1, z1 = pose[3 * joint1], pose[3 * joint1 + 1], pose[3 * joint1 + 2]
        x2, y2, z2 = pose[3 * joint2], pose[3 * joint2 + 1], pose[3 * joint2 + 2]

        # Draw the line between the two joints (bone)
        ax.plot(
            [x1 + trans[0], x2 + trans[0]],
            [y1 + trans[1], y2 + trans[1]],
            [z1 + trans[2], z2 + trans[2]],
            color='blue'
        )

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

def animate_skeleton(pose_world_data, trans_world_data):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    ani = animation.FuncAnimation(
        fig, update_skeleton_frame, frames=len(pose_world_data), fargs=(pose_world_data, trans_world_data, ax),
        interval=100  # Adjust the interval for speed
    )

    plt.show()

# Animate the skeleton
animate_skeleton(pose_world, trans_world)

Summary:

Bones are drawn between joints using matplotlib’s ax.plot() function to represent the skeleton in 3D.
The skeleton follows the SMPL model’s 24 joints, and the connections between them are based on the SMPL bone structure.
You can visualize individual frames or animate the skeleton over time with the body pose data from pose_world and translations from trans_world.

This should give you a complete visualization of the SMPL skeleton and allow you to inspect or animate the body pose frame-by-frame.

The Human3.6M (H36M) dataset uses a different set of joints and connections compared to the SMPL model. In Human3.6M, the skeleton consists of 32 joints, but commonly only 17 joints are used in most 3D pose estimation tasks (similar to other pose estimation datasets).

Here are the 17 main joints of the Human3.6M body model:

0: Pelvis (Hip)
1: Left Hip
2: Left Knee
3: Left Ankle
4: Right Hip
5: Right Knee
6: Right Ankle
7: Spine
8: Thorax (Upper Spine)
9: Neck
10: Head
11: Left Shoulder
12: Left Elbow
13: Left Wrist
14: Right Shoulder
15: Right Elbow
16: Right Wrist

Joint Connections (Human3.6M Skeleton)

The following are the joint connections that define the bones of the Human3.6M skeleton:

h36m_skeleton = [
    (0, 1), (1, 2), (2, 3),    # Pelvis to left ankle
    (0, 4), (4, 5), (5, 6),    # Pelvis to right ankle
    (0, 7), (7, 8),            # Pelvis to thorax (spine)
    (8, 9), (9, 10),           # Thorax to head
    (8, 11), (11, 12), (12, 13), # Left shoulder to left wrist
    (8, 14), (14, 15), (15, 16)  # Right shoulder to right wrist
]

Visualization Example

You can use the same approach as we used for the SMPL model, but with the Human3.6M joint connections. Here’s an example:

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

# Human3.6M skeleton connections
h36m_skeleton = [
    (0, 1), (1, 2), (2, 3),    # Pelvis to left ankle
    (0, 4), (4, 5), (5, 6),    # Pelvis to right ankle
    (0, 7), (7, 8),            # Pelvis to thorax (spine)
    (8, 9), (9, 10),           # Thorax to head
    (8, 11), (11, 12), (12, 13), # Left shoulder to left wrist
    (8, 14), (14, 15), (15, 16)  # Right shoulder to right wrist
]

# Visualize a single frame with the skeleton drawn as bones (Human3.6M model)
def visualize_h36m_skeleton_frame(pose, trans):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # Plot the root translation (as a point)
    ax.scatter(trans[0], trans[1], trans[2], c='r', marker='o', label="Root Translation")

    # Plot each bone in the skeleton
    for joint1, joint2 in h36m_skeleton:
        # Get the 3D coordinates of the two joints
        x1, y1, z1 = pose[3 * joint1], pose[3 * joint1 + 1], pose[3 * joint1 + 2]
        x2, y2, z2 = pose[3 * joint2], pose[3 * joint2 + 1], pose[3 * joint2 + 2]

        # Draw the line between the two joints (bone)
        ax.plot(
            [x1 + trans[0], x2 + trans[0]],
            [y1 + trans[1], y2 + trans[1]],
            [z1 + trans[2], z2 + trans[2]],
            color='blue'
        )

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    plt.legend()
    plt.show()

# Example: visualize the first frame (assuming pose_world and trans_world are for Human3.6M)
visualize_h36m_skeleton_frame(pose_world[0], trans_world[0])

Explanation:

Human3.6M Joint Connections:
- The joint connections (h36m_skeleton) are defined based on the Human3.6M 17-joint body model.
- These connections describe how the joints are linked to form the human skeleton (e.g., pelvis to knee, knee to ankle, etc.).
Visualizing Bones:
- For each pair of connected joints (joint1, joint2), a line is drawn between the two joint positions in 3D space to represent the bone.
Translation:
- The root translation (trans_world) is applied to all joints to move the skeleton into world coordinates.

Animation:

You can animate the Human3.6M skeleton in a similar fashion by adapting the animation code previously provided. Here’s a quick example of how you could animate the Human3.6M skeleton:

import matplotlib.animation as animation

def update_h36m_skeleton_frame(num, pose_world_data, trans_world_data, ax):
    ax.clear()

    pose = pose_world_data[num]
    trans = trans_world_data[num]

    # Plot each bone in the skeleton
    for joint1, joint2 in h36m_skeleton:
        x1, y1, z1 = pose[3 * joint1], pose[3 * joint1 + 1], pose[3 * joint1 + 2]
        x2, y2, z2 = pose[3 * joint2], pose[3 * joint2 + 1], pose[3 * joint2 + 2]

        ax.plot(
            [x1 + trans[0], x2 + trans[0]],
            [y1 + trans[1], y2 + trans[1]],
            [z1 + trans[2], z2 + trans[2]],
            color='blue'
        )

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

def animate_h36m_skeleton(pose_world_data, trans_world_data):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    ani = animation.FuncAnimation(
        fig, update_h36m_skeleton_frame, frames=len(pose_world_data), fargs=(pose_world_data, trans_world_data, ax),
        interval=100  # Adjust the interval for speed
    )

    plt.show()

# Animate the skeleton (assuming pose_world and trans_world are for Human3.6M)
animate_h36m_skeleton(pose_world, trans_world)

Summary:

The Human3.6M skeleton consists of 17 main joints with predefined connections (bones) that define how the joints are linked.
The code above visualizes the skeleton in 3D, with lines drawn between connected joints to represent bones.
You can also animate the skeleton frame-by-frame if you have pose and translation data for multiple frames.

This should give you a good visual representation of the Human3.6M body model and allow you to animate or inspect individual frames from your dataset.

To convert the SMPL joint rotations into actual 3D joint positions (points), we need to follow a few steps. The SMPL model provides joint rotations in axis-angle format, but to get the 3D positions of the joints (which we can visualize as points in 3D space), we need to:

Apply Forward Kinematics: Use the axis-angle rotations of each joint and its parent to calculate the final joint positions.
Hierarchical Structure: SMPL defines the body joints in a hierarchical way, where each joint is positioned relative to its parent. To get the global position of each joint, we need to transform the local rotations and translations through the hierarchy.
Root Joint: Typically, the first joint (pelvis) is the root, and its position is given directly. The rest of the joints are relative to this root.
Convert Axis-Angle to Rotation Matrices: We need to convert the axis-angle representations (3D vectors for each joint) into rotation matrices, which can then be applied to compute joint transformations.
Accumulate Joint Transformations: For each joint, apply the rotation and translation from its parent joint to compute its global position.

Overview of Steps:

Convert axis-angle rotations to rotation matrices.
Apply forward kinematics through the joint hierarchy to compute joint positions.
Visualize the joints and bones in 3D.

Let’s break this process down step by step.

Step 1: Convert Axis-Angle to Rotation Matrices

SMPL uses axis-angle representation to represent rotations. We can convert this into a rotation matrix using the Rodrigues’ rotation formula.

Here’s how you can do this in Python:

import numpy as np
from scipy.spatial.transform import Rotation as R

# Convert axis-angle vector to a rotation matrix
def axis_angle_to_matrix(axis_angle):
    rotation_matrix = R.from_rotvec(axis_angle).as_matrix()
    return rotation_matrix

Step 2: Apply Forward Kinematics to Compute Joint Positions

We will calculate the global joint positions by applying forward kinematics. For each joint, we compute the transformation by accumulating rotations and translations from the root (pelvis) down to the joint of interest.

Here’s how you can compute the joint positions based on the relative rotations:

def compute_joint_positions(pose, trans, joint_hierarchy, joint_rest_positions):
    """
    Compute the global joint positions based on relative rotations and translations.

    :param pose: (23x3) array of axis-angle rotations for each joint
    :param trans: Translation of the root joint (pelvis)
    :param joint_hierarchy: A list of parent indices for each joint
    :param joint_rest_positions: The rest positions of the joints (3D positions of joints in T-pose)

    :return: 24x3 global joint positions in 3D space
    """
    num_joints = len(joint_hierarchy) + 1  # We have 24 joints
    joint_positions = np.zeros((num_joints, 3))  # Global positions of the joints

    # Start with the root joint (pelvis) at the global translation (trans)
    joint_positions[0] = trans

    # Initialize rotation matrices
    joint_rotations = np.zeros((num_joints, 3, 3))
    joint_rotations[0] = np.eye(3)  # Identity matrix for the root joint

    # Iterate over all the joints and compute their global positions
    for i in range(1, num_joints):
        parent_idx = joint_hierarchy[i - 1]

        # Convert axis-angle to rotation matrix for the current joint
        joint_rot = axis_angle_to_matrix(pose[i - 1])

        # Compute the global rotation by multiplying parent's global rotation with current local rotation
        joint_rotations[i] = joint_rotations[parent_idx] @ joint_rot

        # Compute the global position by applying the parent's rotation and adding the local position
        joint_positions[i] = joint_positions[parent_idx] + joint_rotations[parent_idx] @ joint_rest_positions[i]

    return joint_positions

Step 3: Hierarchical Structure of SMPL Joints

SMPL has a hierarchical joint structure where each joint has a parent. Here’s an example of the parent-child relationships for the SMPL 24-joint model:

joint_hierarchy = [
    0,  # Left Hip (1) -> Pelvis (0)
    0,  # Right Hip (2) -> Pelvis (0)
    0,  # Spine (3) -> Pelvis (0)
    1,  # Left Knee (4) -> Left Hip (1)
    2,  # Right Knee (5) -> Right Hip (2)
    3,  # Spine 2 (6) -> Spine (3)
    4,  # Left Ankle (7) -> Left Knee (4)
    5,  # Right Ankle (8) -> Right Knee (5)
    6,  # Spine 3 (9) -> Spine 2 (6)
    7,  # Left Toe (10) -> Left Ankle (7)
    8,  # Right Toe (11) -> Right Ankle (8)
    9,  # Neck (12) -> Spine 3 (9)
    12, # Head (13) -> Neck (12)
    12, # Left Shoulder (14) -> Neck (12)
    12, # Right Shoulder (15) -> Neck (12)
    14, # Left Elbow (16) -> Left Shoulder (14)
    15, # Right Elbow (17) -> Right Shoulder (15)
    16, # Left Wrist (18) -> Left Elbow (16)
    17, # Right Wrist (19) -> Right Elbow (17)
    18, # Left Hand (20) -> Left Wrist (18)
    19, # Right Hand (21) -> Right Wrist (19)
    13, # Left Ear/Nose (22) -> Head (13)
    13  # Right Ear/Nose (23) -> Head (13)
]

Step 4: Visualize the Joint Positions

You can now visualize the computed joint positions in 3D. Here’s an example of how to visualize the skeleton:

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# Visualize the computed joint positions
def visualize_smpl_skeleton(joint_positions, joint_connections):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # Plot the joints
    ax.scatter(joint_positions[:, 0], joint_positions[:, 1], joint_positions[:, 2], c='r', marker='o', label='Joints')

    # Draw the bones
    for joint1, joint2 in joint_connections:
        ax.plot(
            [joint_positions[joint1, 0], joint_positions[joint2, 0]],
            [joint_positions[joint1, 1], joint_positions[joint2, 1]],
            [joint_positions[joint1, 2], joint_positions[joint2, 2]],
            color='blue'
        )

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    plt.legend()
    plt.show()

# Example joint rest positions (in T-pose)
joint_rest_positions = np.zeros((24, 3))  # Replace with actual joint rest positions (T-pose)

# Compute the joint positions from pose and trans
pose = np.random.randn(23, 3)  # Replace with actual pose data
trans = np.random.randn(3)     # Replace with actual root translation
joint_positions = compute_joint_positions(pose, trans, joint_hierarchy, joint_rest_positions)

# Example joint connections (similar to before)
joint_connections = [
    (0, 1), (0, 2), (0, 3), (1, 4), (2, 5), (3, 6), (4, 7), (5, 8), (6, 9),
    (7, 10), (8, 11), (9, 12), (12, 13), (12, 14), (14, 16), (15, 17), (16, 18),
    (17, 19), (18, 20), (19, 21), (9, 15), (9, 22), (9, 23)
]

# Visualize the skeleton
visualize_smpl_skeleton(joint_positions, joint_connections)

Key Points:

Axis-Angle to Rotation: Convert each axis-angle rotation to a rotation matrix.
Forward Kinematics: Use forward kinematics to propagate joint rotations and compute the global positions of each joint.
Visualize the Skeleton: Visualize the joints and bones in 3D using matplotlib.

Conclusion:

By using forward kinematics, you can compute the global joint positions of the SMPL model given the axis-angle rotations. Once the positions are computed, you can visualize the skeleton in 3D using matplotlib.

this is what happens to the joints. can i maybe modify this to get the joints in 3d space?

    pred_body_pose = matrix_to_axis_angle(pred['poses_body']).cpu().numpy().reshape(-1, 69)
    pred_root = matrix_to_axis_angle(pred['poses_root_cam']).cpu().numpy().reshape(-1, 3)
    pred_root_world = matrix_to_axis_angle(pred['poses_root_world']).cpu().numpy().reshape(-1, 3)
    pred_pose = np.concatenate((pred_root, pred_body_pose), axis=-1)
    pred_pose_world = np.concatenate((pred_root_world, pred_body_pose), axis=-1)
    pred_trans = (pred['trans_cam'] - network.output.offset).cpu().numpy()

Yes, you can modify this code to compute the 3D joint positions based on the provided axis-angle rotations (pred_pose_world and pred_trans) and the SMPL model’s joint hierarchy.

Here’s an outline of how to proceed:

Plan:

Understand the Data:
- pred_pose_world: This is the concatenated axis-angle rotation vector for the root joint and the body pose joints. It contains the relative rotations for all the SMPL joints.
- pred_trans: This is the translation of the root joint (pelvis) in world space. It represents the global position of the root joint in 3D space.
Convert the Axis-Angle Representations into rotation matrices.
- Each joint’s rotation is given in axis-angle format. We need to convert these rotations into 3D rotation matrices.
Apply Forward Kinematics:
- Use forward kinematics through the SMPL hierarchy to compute the global 3D joint positions.
Modify the Current Code:
- Extract the necessary axis-angle rotations and translations.
- Compute the 3D joint positions from these rotations and translations.

Step-by-Step Solution

Let’s start by converting the axis-angle rotations into rotation matrices and applying forward kinematics to get the 3D positions of the joints.

1. Modify Code to Get Joint Positions in 3D

You can modify the existing code like this:

import numpy as np
from scipy.spatial.transform import Rotation as R

# SMPL Joint hierarchy
smpl_joint_hierarchy = [
    0,  # Left Hip -> Pelvis
    0,  # Right Hip -> Pelvis
    0,  # Spine -> Pelvis
    1,  # Left Knee -> Left Hip
    2,  # Right Knee -> Right Hip
    3,  # Spine 2 -> Spine
    4,  # Left Ankle -> Left Knee
    5,  # Right Ankle -> Right Knee
    6,  # Spine 3 -> Spine 2
    7,  # Left Toe -> Left Ankle
    8,  # Right Toe -> Right Ankle
    9,  # Neck -> Spine 3
    12, # Head -> Neck
    12, # Left Shoulder -> Neck
    12, # Right Shoulder -> Neck
    13, # Left Elbow -> Left Shoulder
    14, # Right Elbow -> Right Shoulder
    15, # Left Wrist -> Left Elbow
    16, # Right Wrist -> Right Elbow
    17, # Left Hand -> Left Wrist
    18, # Right Hand -> Right Wrist
    13, # Nose -> Head
    13  # Eyes -> Head
]

# Convert axis-angle to rotation matrix
def axis_angle_to_matrix(axis_angle):
    return R.from_rotvec(axis_angle).as_matrix()

# Compute joint positions using forward kinematics
def compute_joint_positions(pred_pose_world, pred_trans, joint_hierarchy, joint_rest_positions):
    num_joints = len(joint_hierarchy) + 1  # We have 24 joints
    joint_positions = np.zeros((num_joints, 3))  # Global positions of the joints

    # Root joint (Pelvis) position
    joint_positions[0] = pred_trans

    # Rotation matrices for each joint
    joint_rotations = np.zeros((num_joints, 3, 3))
    joint_rotations[0] = np.eye(3)  # Root has no rotation initially

    # Iterate over all the joints to calculate global positions
    for i in range(1, num_joints):
        parent_idx = joint_hierarchy[i - 1]

        # Convert axis-angle to rotation matrix
        local_rotation = axis_angle_to_matrix(pred_pose_world[3 * (i - 1): 3 * i])

        # Compute the global rotation of the joint by applying parent's rotation
        joint_rotations[i] = joint_rotations[parent_idx] @ local_rotation

        # Compute the global position
        joint_positions[i] = joint_positions[parent_idx] + joint_rotations[parent_idx] @ joint_rest_positions[i]

    return joint_positions

# Example joint rest positions (T-pose, replace with actual rest positions)
joint_rest_positions = np.zeros((24, 3))  # Replace with real rest pose values

# Modify this to extract the pred_pose_world and pred_trans from your existing code
pred_pose_world = np.concatenate((pred_root_world, pred_body_pose), axis=-1).reshape(-1, 3)  # Reshape to (23, 3)
pred_trans = (pred['trans_cam'] - network.output.offset).cpu().numpy()  # Root translation

# Compute 3D joint positions
joint_positions = compute_joint_positions(pred_pose_world, pred_trans, smpl_joint_hierarchy, joint_rest_positions)

Explanation of the Code:

axis_angle_to_matrix: Converts the axis-angle rotations to rotation matrices using the scipy.spatial.transform.Rotation class.
compute_joint_positions: This function applies forward kinematics to compute the 3D joint positions:
- For each joint, it computes the global rotation by accumulating rotations from its parent.
- It then computes the global position by transforming the joint’s rest position and adding it to the parent joint’s position.
joint_rest_positions: This is the rest pose (typically the T-pose) of the SMPL model. In the T-pose, each joint has a known rest position relative to its parent. You need to replace the placeholder joint_rest_positions with the actual T-pose values.
pred_pose_world: Extracted from the existing code. The axis-angle rotations for each joint are concatenated and reshaped into (23, 3) where each joint has a 3D rotation vector.
pred_trans: Represents the global translation of the root joint (pelvis), which is extracted from pred['trans_cam'].

Step 4: Visualization

You can use the earlier visualization code to plot the skeleton in 3D after computing the joint positions:

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# Visualize the SMPL skeleton
def visualize_smpl_skeleton(joint_positions, joint_connections):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # Plot the joints
    ax.scatter(joint_positions[:, 0], joint_positions[:, 1], joint_positions[:, 2], c='r', marker='o', label='Joints')

    # Draw the bones
    for joint1, joint2 in joint_connections:
        ax.plot(
            [joint_positions[joint1, 0], joint_positions[joint2, 0]],
            [joint_positions[joint1, 1], joint_positions[joint2, 1]],
            [joint_positions[joint1, 2], joint_positions[joint2, 2]],
            color='blue'
        )

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    plt.legend()
    plt.show()

# Example joint connections
joint_connections = [
    (0, 1), (0, 2), (0, 3), (1, 4), (2, 5), (3, 6), (4, 7), (5, 8), (6, 9),
    (7, 10), (8, 11), (9, 12), (12, 13), (12, 14), (14, 16), (15, 17), (16, 18),
    (17, 19), (18, 20), (19, 21), (9, 15), (9, 22), (9, 23)
]

# Visualize the skeleton
visualize_smpl_skeleton(joint_positions, joint_connections)

Key Points:

Forward Kinematics: The hierarchical structure of the SMPL model is respected, and each joint’s position is computed relative to its parent.
Global Position: The global position of each joint is computed by applying the parent joint’s rotation and adding the rest position.
Visualization: You can visualize the skeleton using matplotlib to plot the joints and the bones.

Conclusion:

You can modify the code to compute the 3D positions of the joints based on the axis-angle rotations and the root translation. By using forward kinematics, you can compute the global positions of all joints and visualize the SMPL skeleton in 3D.

Understood! You want to take the axis-angle representations of the 24 joints from the SMPL skeleton and convert them into 3D (x, y, z) joint positions. This involves:

Applying Forward Kinematics: Use the axis-angle rotations and the known resting pose to calculate the joint positions.
Default Resting Joint Positions: I will provide the standard T-pose used in the SMPL model.

Steps to Solve:

Convert Axis-Angle to Rotation Matrices.
Apply Forward Kinematics to compute the global 3D positions of each joint based on the SMPL hierarchy and the resting positions.
Use the provided SMPL T-pose as the default rest joint positions.

Key Information:

SMPL T-pose: The rest positions of the SMPL joints are known and correspond to a default T-pose.
Forward Kinematics: Given the parent-child relationships in the SMPL hierarchy, we calculate the global 3D positions of each joint by accumulating the transformations from its parent.

SMPL T-Pose (Rest Positions)

Here are the standard SMPL T-pose joint positions in meters. These are the 3D rest positions of the 24 SMPL joints in a T-pose (without any pose deformations).

# SMPL default T-pose joint positions (in meters)
smpl_rest_pose = np.array([
    [0.000000, 0.000000, 0.000000],    # Pelvis (root joint)
    [0.065231, -0.084792, 0.004014],   # Left Hip
    [-0.065231, -0.084792, 0.004014],  # Right Hip
    [0.000000, 0.059053, 0.000000],    # Spine
    [0.091713, -0.370266, 0.004094],   # Left Knee
    [-0.091713, -0.370266, 0.004094],  # Right Knee
    [0.000000, 0.206089, 0.000000],    # Spine 2
    [0.097469, -0.738293, 0.004016],   # Left Ankle
    [-0.097469, -0.738293, 0.004016],  # Right Ankle
    [0.000000, 0.451070, 0.000000],    # Spine 3
    [0.097515, -0.846017, -0.085256],  # Left Toe
    [-0.097515, -0.846017, -0.085256], # Right Toe
    [0.000000, 0.628839, 0.000000],    # Neck
    [0.000000, 0.673400, 0.000000],    # Head
    [0.173363, 0.573362, 0.000052],    # Left Shoulder
    [-0.173363, 0.573362, 0.000052],   # Right Shoulder
    [0.325267, 0.572304, 0.000147],    # Left Elbow
    [-0.325267, 0.572304, 0.000147],   # Right Elbow
    [0.431036, 0.572345, 0.000304],    # Left Wrist
    [-0.431036, 0.572345, 0.000304],   # Right Wrist
    [0.498760, 0.571109, 0.000320],    # Left Hand
    [-0.498760, 0.571109, 0.000320],   # Right Hand
    [0.000000, 0.719092, 0.025512],    # Nose
    [0.000000, 0.731636, 0.037334],    # Eyes
])

SMPL Joint Hierarchy

The parent-child relationship of the SMPL joints can be described as follows. Each joint has a parent, and its position is relative to that parent.

# SMPL joint hierarchy (indices of parents)
smpl_joint_hierarchy = [
    0,  # Left Hip -> Pelvis
    0,  # Right Hip -> Pelvis
    0,  # Spine -> Pelvis
    1,  # Left Knee -> Left Hip
    2,  # Right Knee -> Right Hip
    3,  # Spine 2 -> Spine
    4,  # Left Ankle -> Left Knee
    5,  # Right Ankle -> Right Knee
    6,  # Spine 3 -> Spine 2
    7,  # Left Toe -> Left Ankle
    8,  # Right Toe -> Right Ankle
    9,  # Neck -> Spine 3
    12, # Head -> Neck
    12, # Left Shoulder -> Neck
    12, # Right Shoulder -> Neck
    13, # Left Elbow -> Left Shoulder
    14, # Right Elbow -> Right Shoulder
    15, # Left Wrist -> Left Elbow
    16, # Right Wrist -> Right Elbow
    17, # Left Hand -> Left Wrist
    18, # Right Hand -> Right Wrist
    13, # Nose -> Head
    13  # Eyes -> Head
]

Convert Axis-Angle to Rotation Matrix

We will first convert the axis-angle representation to a rotation matrix using scipy.

import numpy as np
from scipy.spatial.transform import Rotation as R

# Convert axis-angle vector to rotation matrix
def axis_angle_to_matrix(axis_angle):
    rotation_matrix = R.from_rotvec(axis_angle).as_matrix()
    return rotation_matrix

Compute 3D Joint Positions Using Forward Kinematics

We will compute the 3D joint positions by applying forward kinematics through the SMPL hierarchy, starting from the pelvis and accumulating the transformations.

def compute_smpl_joint_positions(pose, trans, joint_hierarchy, joint_rest_positions):
    """
    Compute the global joint positions in 3D based on SMPL model pose and rest positions.

    :param pose: (23x3) array of axis-angle rotations for each joint
    :param trans: (3) translation of the root joint (pelvis) in world space
    :param joint_hierarchy: List of parent indices for each joint
    :param joint_rest_positions: (24x3) rest positions of the joints (T-pose)

    :return: (24x3) global joint positions in 3D
    """
    num_joints = len(joint_hierarchy) + 1  # We have 24 joints (Pelvis + 23 others)
    joint_positions = np.zeros((num_joints, 3))  # To store the global positions of each joint
    joint_rotations = np.zeros((num_joints, 3, 3))  # To store the global rotation matrices for each joint

    # Start with the root joint (Pelvis)
    joint_positions[0] = trans
    joint_rotations[0] = np.eye(3)  # The pelvis has no rotation initially

    # Iterate over each joint (starting from index 1) and compute its position
    for i in range(1, num_joints):
        parent_idx = joint_hierarchy[i - 1]

        # Convert axis-angle to rotation matrix for the current joint
        joint_rotation_matrix = axis_angle_to_matrix(pose[i - 1])

        # Compute the global rotation for the current joint
        joint_rotations[i] = joint_rotations[parent_idx] @ joint_rotation_matrix

        # Compute the global position for the current joint
        joint_positions[i] = joint_positions[parent_idx] + joint_rotations[parent_idx] @ joint_rest_positions[i]

    return joint_positions

Usage Example

Let’s now compute the 3D joint positions using the pose (axis-angle rotations) and the root joint translation (trans). Here’s how to use the above code with some example data:

# Example: Axis-angle pose data (23 joints x 3)
pose = np.random.randn(23, 3)  # Replace with actual pose data from SMPL
trans = np.array([0.0, 0.0, 0.9])  # Example root joint (pelvis) translation in meters

# Compute the 3D joint positions using the T-pose rest positions
joint_positions = compute_smpl_joint_positions(pose, trans, smpl_joint_hierarchy, smpl_rest_pose)

# Print the resulting joint positions
print("Joint positions in 3D (meters):")
print(joint_positions)

Visualization (Optional)

You can visualize the resulting 3D joint positions as a skeleton using matplotlib:

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def visualize_smpl_skeleton(joint_positions, joint_connections):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # Plot the joints
    ax.scatter(joint_positions[:, 0], joint_positions[:, 1], joint_positions[:, 2], c='r', marker='o', label='Joints')

    # Draw the bones
    for joint1, joint2 in joint_connections:
        ax.plot(
            [joint_positions[joint1, 0], joint_positions[joint2, 0]],
            [joint_positions[joint1, 1], joint_positions[joint2, 1]],
            [joint_positions[joint1, 2], joint_positions[joint2, 2]],
            color='blue'
        )

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    plt.legend()
    plt.show()

# Example joint connections (same as before)
joint_connections = [
    (0, 1), (0, 2), (0, 3), (1, 4), (2, 5), (3, 6), (4, 7), (5, 8), (6, 9),
    (7, 10), (8, 11), (9, 12), (12, 13), (12, 14), (14, 16), (15, 17), (16, 18),
    (17, 19), (18, 20), (19, 21), (9, 15), (9, 22), (9, 23)
]

# Visualize the skeleton
visualize_smpl_skeleton(joint_positions, joint_connections)

Conclusion:

This code takes the SMPL 24 joint skeleton in axis-angle representation and computes the global 3D joint positions using forward kinematics. The T-pose rest joint positions provide a default skeleton, and the rotations from the pose data deform this skeleton to represent a new pose in 3D space.

You can now use this to compute the 3D joint positions from any SMPL pose vector and visualize the results.

If your axis-angle data has a shape of (72,), it means that the data represents 24 joints (with 3 axis-angle values for each joint). Specifically, 72 values corresponds to 24 joints × 3 values per joint (representing the axis-angle rotation for each joint).

To accommodate this format, you need to reshape the (72,) axis-angle data into (24, 3), where each row corresponds to the 3D axis-angle rotation for one joint.

Step-by-Step Adjustments

Reshape the axis-angle data from (72,) to (24, 3), where each joint has its own 3D rotation vector.
Ensure that the first joint is the root joint (Pelvis), which is typically represented in the translation (pred_trans).
Apply forward kinematics using the reshaped axis-angle data.

Adjusted Code for the (72,) Axis-Angle Data

We’ll modify the earlier example to handle the (72,)-shaped axis-angle data.

Updated Code

import numpy as np
from scipy.spatial.transform import Rotation as R

# Convert axis-angle vector to rotation matrix
def axis_angle_to_matrix(axis_angle):
    return R.from_rotvec(axis_angle).as_matrix()

def compute_smpl_joint_positions(pose, trans, joint_hierarchy, joint_rest_positions):
    """
    Compute the global joint positions in 3D based on SMPL model pose and rest positions.

    :param pose: (24x3) array of axis-angle rotations for each joint
    :param trans: (3) translation of the root joint (pelvis) in world space
    :param joint_hierarchy: List of parent indices for each joint
    :param joint_rest_positions: (24x3) rest positions of the joints (T-pose)

    :return: (24x3) global joint positions in 3D
    """
    num_joints = len(joint_hierarchy) + 1  # We have 24 joints (Pelvis + 23 others)
    joint_positions = np.zeros((num_joints, 3))  # To store the global positions of each joint
    joint_rotations = np.zeros((num_joints, 3, 3))  # To store the global rotation matrices for each joint

    # Start with the root joint (Pelvis)
    joint_positions[0] = trans
    joint_rotations[0] = np.eye(3)  # The pelvis has no rotation initially

    # Iterate over each joint (starting from index 1) and compute its position
    for i in range(1, num_joints):
        parent_idx = joint_hierarchy[i - 1]

        # Convert axis-angle to rotation matrix for the current joint
        joint_rotation_matrix = axis_angle_to_matrix(pose[i])

        # Compute the global rotation for the current joint
        joint_rotations[i] = joint_rotations[parent_idx] @ joint_rotation_matrix

        # Compute the global position for the current joint
        joint_positions[i] = joint_positions[parent_idx] + joint_rotations[parent_idx] @ joint_rest_positions[i]

    return joint_positions

# SMPL Joint hierarchy
smpl_joint_hierarchy = [
    0,  # Left Hip -> Pelvis
    0,  # Right Hip -> Pelvis
    0,  # Spine -> Pelvis
    1,  # Left Knee -> Left Hip
    2,  # Right Knee -> Right Hip
    3,  # Spine 2 -> Spine
    4,  # Left Ankle -> Left Knee
    5,  # Right Ankle -> Right Knee
    6,  # Spine 3 -> Spine 2
    7,  # Left Toe -> Left Ankle
    8,  # Right Toe -> Right Ankle
    9,  # Neck -> Spine 3
    12, # Head -> Neck
    12, # Left Shoulder -> Neck
    12, # Right Shoulder -> Neck
    13, # Left Elbow -> Left Shoulder
    14, # Right Elbow -> Right Shoulder
    15, # Left Wrist -> Left Elbow
    16, # Right Wrist -> Right Elbow
    17, # Left Hand -> Left Wrist
    18, # Right Hand -> Right Wrist
    13, # Nose -> Head
    13  # Eyes -> Head
]

# SMPL default T-pose joint positions (in meters)
smpl_rest_pose = np.array([
    [0.000000, 0.000000, 0.000000],    # Pelvis (root joint)
    [0.065231, -0.084792, 0.004014],   # Left Hip
    [-0.065231, -0.084792, 0.004014],  # Right Hip
    [0.000000, 0.059053, 0.000000],    # Spine
    [0.091713, -0.370266, 0.004094],   # Left Knee
    [-0.091713, -0.370266, 0.004094],  # Right Knee
    [0.000000, 0.206089, 0.000000],    # Spine 2
    [0.097469, -0.738293, 0.004016],   # Left Ankle
    [-0.097469, -0.738293, 0.004016],  # Right Ankle
    [0.000000, 0.451070, 0.000000],    # Spine 3
    [0.097515, -0.846017, -0.085256],  # Left Toe
    [-0.097515, -0.846017, -0.085256], # Right Toe
    [0.000000, 0.628839, 0.000000],    # Neck
    [0.000000, 0.673400, 0.000000],    # Head
    [0.173363, 0.573362, 0.000052],    # Left Shoulder
    [-0.173363, 0.573362, 0.000052],   # Right Shoulder
    [0.325267, 0.572304, 0.000147],    # Left Elbow
    [-0.325267, 0.572304, 0.000147],   # Right Elbow
    [0.431036, 0.572345, 0.000304],    # Left Wrist
    [-0.431036, 0.572345, 0.000304],   # Right Wrist
    [0.498760, 0.571109, 0.000320],    # Left Hand
    [-0.498760, 0.571109, 0.000320],   # Right Hand
    [0.000000, 0.719092, 0.025512],    # Nose
    [0.000000, 0.731636, 0.037334],    # Eyes
])

# Example axis-angle pose data (72 values, reshape to 24x3)
pose = np.random.randn(72).reshape(24, 3)  # Replace with actual axis-angle data
trans = np.array([0.0, 0.0, 0.9])  # Example root joint translation (pelvis) in meters

# Compute 3D joint positions
joint_positions = compute_smpl_joint_positions(pose, trans, smpl_joint_hierarchy, smpl_rest_pose)

# Print the resulting joint positions
print("Joint positions in 3D (meters):")
print(joint_positions)

Key Adjustments:

Reshape Axis-Angle Data: The pose data is initially in the shape (72,), so we reshape it to (24, 3) to represent the axis-angle rotations for all 24 joints:
```
pose = np.random.randn(72).reshape(24, 3)  # Replace with actual axis-angle data
```
Root Joint (Pelvis): The root joint (Pelvis) translation is still given by trans. This represents the global position of the root joint, and we apply forward kinematics to the other joints relative to the pelvis.

Visualization (Optional):

You can visualize the 3D joint positions using the visualize_smpl_skeleton function provided earlier. This will show the skeleton in 3D with lines connecting the joints.

Conclusion:

By reshaping the (72,) axis-angle data into (24, 3), you can compute the 3D joint positions of the SMPL skeleton using forward kinematics. The code provided computes and visualizes the joints in 3D based on the given pose and translation data.

If your data contains multiple frames of an animation, the axis-angle data will have a 3D shape instead of a 2D or 1D shape. Specifically, the shape could be something like (num_frames, 72), where:

num_frames: Number of frames in the animation.
72: The axis-angle data for all 24 joints, with 3 values (axis-angle) per joint.

In this case, you would need to:

Reshape the data per frame to compute the 3D joint positions for each frame.
Loop through each frame, applying forward kinematics to compute the global 3D joint positions for all the frames.
Visualize or store the computed joint positions frame by frame.

Plan:

Reshape the Axis-Angle Data: Reshape the data into (num_frames, 24, 3) to process the axis-angle rotations for each frame.
Iterate Over Each Frame: For each frame, use forward kinematics to compute the 3D joint positions.
Store or Visualize the Results.

Code to Handle 3D Shape (Multiple Frames)

Let’s modify the earlier code to handle multiple frames of axis-angle data.

Adjusted Code for Multiple Frames

import numpy as np
from scipy.spatial.transform import Rotation as R

# Convert axis-angle vector to rotation matrix
def axis_angle_to_matrix(axis_angle):
    return R.from_rotvec(axis_angle).as_matrix()

def compute_smpl_joint_positions(pose, trans, joint_hierarchy, joint_rest_positions):
    """
    Compute the global joint positions in 3D based on SMPL model pose and rest positions.

    :param pose: (24x3) array of axis-angle rotations for each joint
    :param trans: (3) translation of the root joint (pelvis) in world space
    :param joint_hierarchy: List of parent indices for each joint
    :param joint_rest_positions: (24x3) rest positions of the joints (T-pose)

    :return: (24x3) global joint positions in 3D
    """
    num_joints = len(joint_hierarchy) + 1  # We have 24 joints (Pelvis + 23 others)
    joint_positions = np.zeros((num_joints, 3))  # To store the global positions of each joint
    joint_rotations = np.zeros((num_joints, 3, 3))  # To store the global rotation matrices for each joint

    # Start with the root joint (Pelvis)
    joint_positions[0] = trans
    joint_rotations[0] = np.eye(3)  # The pelvis has no rotation initially

    # Iterate over each joint (starting from index 1) and compute its position
    for i in range(1, num_joints):
        parent_idx = joint_hierarchy[i - 1]

        # Convert axis-angle to rotation matrix for the current joint
        joint_rotation_matrix = axis_angle_to_matrix(pose[i])

        # Compute the global rotation for the current joint
        joint_rotations[i] = joint_rotations[parent_idx] @ joint_rotation_matrix

        # Compute the global position for the current joint
        joint_positions[i] = joint_positions[parent_idx] + joint_rotations[parent_idx] @ joint_rest_positions[i]

    return joint_positions

# SMPL Joint hierarchy
smpl_joint_hierarchy = [
    0,  # Left Hip -> Pelvis
    0,  # Right Hip -> Pelvis
    0,  # Spine -> Pelvis
    1,  # Left Knee -> Left Hip
    2,  # Right Knee -> Right Hip
    3,  # Spine 2 -> Spine
    4,  # Left Ankle -> Left Knee
    5,  # Right Ankle -> Right Knee
    6,  # Spine 3 -> Spine 2
    7,  # Left Toe -> Left Ankle
    8,  # Right Toe -> Right Ankle
    9,  # Neck -> Spine 3
    12, # Head -> Neck
    12, # Left Shoulder -> Neck
    12, # Right Shoulder -> Neck
    13, # Left Elbow -> Left Shoulder
    14, # Right Elbow -> Right Shoulder
    15, # Left Wrist -> Left Elbow
    16, # Right Wrist -> Right Elbow
    17, # Left Hand -> Left Wrist
    18, # Right Hand -> Right Wrist
    13, # Nose -> Head
    13  # Eyes -> Head
]

# SMPL default T-pose joint positions (in meters)
smpl_rest_pose = np.array([
    [0.000000, 0.000000, 0.000000],    # Pelvis (root joint)
    [0.065231, -0.084792, 0.004014],   # Left Hip
    [-0.065231, -0.084792, 0.004014],  # Right Hip
    [0.000000, 0.059053, 0.000000],    # Spine
    [0.091713, -0.370266, 0.004094],   # Left Knee
    [-0.091713, -0.370266, 0.004094],  # Right Knee
    [0.000000, 0.206089, 0.000000],    # Spine 2
    [0.097469, -0.738293, 0.004016],   # Left Ankle
    [-0.097469, -0.738293, 0.004016],  # Right Ankle
    [0.000000, 0.451070, 0.000000],    # Spine 3
    [0.097515, -0.846017, -0.085256],  # Left Toe
    [-0.097515, -0.846017, -0.085256], # Right Toe
    [0.000000, 0.628839, 0.000000],    # Neck
    [0.000000, 0.673400, 0.000000],    # Head
    [0.173363, 0.573362, 0.000052],    # Left Shoulder
    [-0.173363, 0.573362, 0.000052],   # Right Shoulder
    [0.325267, 0.572304, 0.000147],    # Left Elbow
    [-0.325267, 0.572304, 0.000147],   # Right Elbow
    [0.431036, 0.572345, 0.000304],    # Left Wrist
    [-0.431036, 0.572345, 0.000304],   # Right Wrist
    [0.498760, 0.571109, 0.000320],    # Left Hand
    [-0.498760, 0.571109, 0.000320],   # Right Hand
    [0.000000, 0.719092, 0.025512],    # Nose
    [0.000000, 0.731636, 0.037334],    # Eyes
])

# Assume pose_data is in the shape (num_frames, 72) -> (num_frames, 24, 3)
# Replace this with actual data
num_frames = 100
pose_data = np.random.randn(num_frames, 72).reshape(num_frames, 24, 3)  # Example pose data for all frames
trans_data = np.random.randn(num_frames, 3)  # Example translation data for all frames

# Iterate over each frame and compute joint positions
all_joint_positions = []

for frame_idx in range(num_frames):
    pose = pose_data[frame_idx]  # Pose for the current frame (24x3)
    trans = trans_data[frame_idx]  # Translation for the current frame (3)

    # Compute the joint positions for the current frame
    joint_positions = compute_smpl_joint_positions(pose, trans, smpl_joint_hierarchy, smpl_rest_pose)

    all_joint_positions.append(joint_positions)

# Convert the list to an array (num_frames, 24, 3)
all_joint_positions = np.array(all_joint_positions)

# Print the joint positions for the first frame
print("Joint positions for the first frame:")
print(all_joint_positions[0])

Explanation:

Handling Multiple Frames:
- pose_data has a shape of (num_frames, 72), which corresponds to num_frames of axis-angle data for all 24 joints.
- The pose data is reshaped to (num_frames, 24, 3) to handle the 24 joints and their 3D axis-angle representation.
Loop Through Frames:
- For each frame, we extract the pose and translation data and compute the joint positions using forward kinematics.
- The joint positions for each frame are stored in all_joint_positions, which will have a final shape of (num_frames, 24, 3).
Final Output:
- After looping through all frames, all_joint_positions contains the 3D joint

You’re right to point out that we need to account for camera transformations and how the data is transformed into world space, particularly in cases where we have pose_world and trans_world. Let’s break this down and make sure we’re handling all transformations properly, especially considering camera transformations and their impact on the resulting 3D joint positions.

Key Data:

pred_pose: The pose of the SMPL skeleton, likely in camera space.
pred_trans: The translation of the root joint (Pelvis) in camera space.
pose_world: The global (world) pose, which includes the joint rotations in world coordinates.
trans_world: The translation in world coordinates for the root joint (Pelvis).

From what you’re describing, you are provided both camera-space data (pred_pose, pred_trans) and world-space data (pose_world, trans_world). Since you want to compute the joint positions in 3D world space, we’ll focus on using pose_world and trans_world, which are already in the global frame.

Key Points:

pose_world contains the axis-angle rotations in world space for each joint.
trans_world provides the root translation in world space, indicating the global position of the pelvis.

To compute the 3D joint positions in world coordinates, we can directly use pose_world and trans_world in a similar manner to what we were doing before but without needing to account for camera transformations (since pose_world and trans_world are already in the global frame).

Updated Code to Handle World Space Data

Here’s how you would adjust the code to use pose_world and trans_world for multiple frames of data:

import numpy as np
from scipy.spatial.transform import Rotation as R

# Convert axis-angle vector to rotation matrix
def axis_angle_to_matrix(axis_angle):
    return R.from_rotvec(axis_angle).as_matrix()

def compute_smpl_joint_positions(pose_world, trans_world, joint_hierarchy, joint_rest_positions):
    """
    Compute the global joint positions in 3D based on SMPL model pose and rest positions.

    :param pose_world: (24x3) array of axis-angle rotations for each joint in world space
    :param trans_world: (3) translation of the root joint (pelvis) in world space
    :param joint_hierarchy: List of parent indices for each joint
    :param joint_rest_positions: (24x3) rest positions of the joints (T-pose)

    :return: (24x3) global joint positions in 3D world space
    """
    num_joints = len(joint_hierarchy) + 1  # We have 24 joints (Pelvis + 23 others)
    joint_positions = np.zeros((num_joints, 3))  # To store the global positions of each joint
    joint_rotations = np.zeros((num_joints, 3, 3))  # To store the global rotation matrices for each joint

    # Start with the root joint (Pelvis)
    joint_positions[0] = trans_world
    joint_rotations[0] = np.eye(3)  # The pelvis has no rotation initially

    # Iterate over each joint (starting from index 1) and compute its position
    for i in range(1, num_joints):
        parent_idx = joint_hierarchy[i - 1]

        # Convert axis-angle to rotation matrix for the current joint (already in world space)
        joint_rotation_matrix = axis_angle_to_matrix(pose_world[i])

        # Compute the global rotation for the current joint
        joint_rotations[i] = joint_rotations[parent_idx] @ joint_rotation_matrix

        # Compute the global position for the current joint
        joint_positions[i] = joint_positions[parent_idx] + joint_rotations[parent_idx] @ joint_rest_positions[i]

    return joint_positions

# SMPL Joint hierarchy (same as before)
smpl_joint_hierarchy = [
    0,  # Left Hip -> Pelvis
    0,  # Right Hip -> Pelvis
    0,  # Spine -> Pelvis
    1,  # Left Knee -> Left Hip
    2,  # Right Knee -> Right Hip
    3,  # Spine 2 -> Spine
    4,  # Left Ankle -> Left Knee
    5,  # Right Ankle -> Right Knee
    6,  # Spine 3 -> Spine 2
    7,  # Left Toe -> Left Ankle
    8,  # Right Toe -> Right Ankle
    9,  # Neck -> Spine 3
    12, # Head -> Neck
    12, # Left Shoulder -> Neck
    12, # Right Shoulder -> Neck
    13, # Left Elbow -> Left Shoulder
    14, # Right Elbow -> Right Shoulder
    15, # Left Wrist -> Left Elbow
    16, # Right Wrist -> Right Elbow
    17, # Left Hand -> Left Wrist
    18, # Right Hand -> Right Wrist
    13, # Nose -> Head
    13  # Eyes -> Head
]

# SMPL default T-pose joint positions (in meters)
smpl_rest_pose = np.array([
    [0.000000, 0.000000, 0.000000],    # Pelvis (root joint)
    [0.065231, -0.084792, 0.004014],   # Left Hip
    [-0.065231, -0.084792, 0.004014],  # Right Hip
    [0.000000, 0.059053, 0.000000],    # Spine
    [0.091713, -0.370266, 0.004094],   # Left Knee
    [-0.091713, -0.370266, 0.004094],  # Right Knee
    [0.000000, 0.206089, 0.000000],    # Spine 2
    [0.097469, -0.738293, 0.004016],   # Left Ankle
    [-0.097469, -0.738293, 0.004016],  # Right Ankle
    [0.000000, 0.451070, 0.000000],    # Spine 3
    [0.097515, -0.846017, -0.085256],  # Left Toe
    [-0.097515, -0.846017, -0.085256], # Right Toe
    [0.000000, 0.628839, 0.000000],    # Neck
    [0.000000, 0.673400, 0.000000],    # Head
    [0.173363, 0.573362, 0.000052],    # Left Shoulder
    [-0.173363, 0.573362, 0.000052],   # Right Shoulder
    [0.325267, 0.572304, 0.000147],    # Left Elbow
    [-0.325267, 0.572304, 0.000147],   # Right Elbow
    [0.431036, 0.572345, 0.000304],    # Left Wrist
    [-0.431036, 0.572345, 0.000304],   # Right Wrist
    [0.498760, 0.571109, 0.000320],    # Left Hand
    [-0.498760, 0.571109, 0.000320],   # Right Hand
    [0.000000, 0.719092, 0.025512],    # Nose
    [0.000000, 0.731636, 0.037334],    # Eyes
])

# Assume pose_world_data and trans_world_data are in shape (num_frames, 24, 3) and (num_frames, 3), respectively
# Replace with actual data
num_frames = 100
pose_world_data = np.random.randn(num_frames, 24, 3)  # Example pose_world data for all frames
trans_world_data = np.random.randn(num_frames, 3)  # Example trans_world data for all frames

# Iterate over each frame and compute joint positions
all_joint_positions = []

for frame_idx in range(num_frames):
    pose_world = pose_world_data[frame_idx]  # Pose for the current frame (24x3 in world space)
    trans_world = trans_world_data[frame_idx]  # Translation for the current frame (3 in world space)

    # Compute the joint positions for the current frame
    joint_positions = compute_smpl_joint_positions(pose_world, trans_world, smpl_joint_hierarchy, smpl_rest_pose)

    all_joint_positions.append(joint_positions)

# Convert the list to an array (num_frames, 24, 3)
all_joint_positions = np.array(all_joint_positions)

# Print the joint positions for the first frame
print("Joint positions for the first frame (in world space):")
print(all_joint_positions[0])

Explanation:

Using pose_world and trans_world:
- We assume that pose_world_data contains the axis-angle rotations for the 24 joints in world space, and trans_world_data contains the translations for the root joint (Pelvis) in world space.
- For each frame, we extract pose_world (shape: (24, 3)) and trans_world (shape: (3,)) and pass them into the compute_smpl_joint_positions() function to compute the 3D joint positions.
Looping Over Frames:
- For each frame, we compute the 3D joint positions and store them in all_joint_positions.
- After looping through all frames, all_joint_positions will have the shape (num_frames, 24, 3), where each frame contains the 3D joint positions of all 24 joints.
Global (World) Coordinates:
- Since we’re using pose_world and trans_world, all computed joint positions are in world space, and no further camera transformations are necessary.

Visualization:

You can visualize the resulting 3D joint positions in world space using the same visualization approach as before.

Conclusion:

This solution computes the 3D joint positions of the SMPL skeleton for multiple frames using the world-space data (pose_world and trans_world). The axis-angle rotations are converted into 3D joint positions, and forward kinematics is applied to propagate the transformations through the joint hierarchy.