skeleton.py

"""Functions for visualizing skeleton movement generated by the RNN"""
from IPython.display import FileLink
from matplotlib.animation import FuncAnimation
import numpy as np
import matplotlib.pyplot as plt
import transforms3d 

def rotate_translate_frame(points, rotation, translation_hrz):
  # translate to put back at origin and rotate to be in the local frame
  points_rot_transl = rotation.dot(points - translation_hrz)
  return points_rot_transl

def inv_rotate_translate_frame(points, rotation, translation_hrz):
  points_rot_transl = rotation.T.dot(points) + translation_hrz  
  return points_rot_transl

def angle2Coordinate(data, data_ini, c_joints, norm_ij):
    dec_delta = data[:, 0:1]
    dec_h = data[:, 1:2]
    dec_d1d2 = data[:, 2:4]
    dec_Sxyz_ij = data[:, 4:]
    dec_Sxyz_ij = dec_Sxyz_ij.reshape(dec_delta.shape[0], -1, 3)

    d1d2_ini = data_ini[0]
    delta_ini = data_ini[1]

    # start with initial condition
    rotation = np.eye(3)
    hrz_transl = np.zeros((3, 1))

    t = 0
    joint_coordinate_arr = []
    delta_cum = 0.
    hrz_transl_cum = np.zeros((3))
    for t in [i for i in range(data.shape[0]-1)]:

        # compute rotation and hrz_transl (need to integrate info for tranzl)
        # hrz_transl = hrz_transl + dec_xyz_root[t, :].reshape(3, 1)

        # compose previous rotation with the new one (need to be extracted from the hips...)
        # extract the rotation from the hips



        if t == 0:
            # rotation = np.eye(3)
            # the first frame is at null rotation and translation
            hrz_transl = np.zeros((3))
            rotation = np.eye(3)
            delta_cum += delta_ini[0]
            hrz_transl = np.dot(rotation.T,np.array([d1d2_ini[0], 0, d1d2_ini[1]]))
            hrz_transl_cum += hrz_transl

            rotation = transforms3d.euler.euler2mat(0,delta_cum,0)
        else:
            delta_cum += dec_delta[t-1, 0]
            rotation = transforms3d.euler.euler2mat(0,delta_cum,0)
            hrz_transl = np.dot(rotation.T,[dec_d1d2[t-1][0], 0, dec_d1d2[t-1][1]])
            hrz_transl_cum += hrz_transl
        
        # print('delta_cum', delta_cum, 'hrz_transl_cum', hrz_transl_cum)
        
        joint_coordinate = dict()
        joint_coordinate_global = dict()
        # joint_coordinate['root'] = inv_rotate_translate_frame(np.array([0, dec_h[t][0], 0]), rotation, hrz_transl_cum)
        joint_coordinate['root'] = np.array([0, dec_h[t][0], 0])
        joint_coordinate_global['root'] = inv_rotate_translate_frame(np.array([0, dec_h[t][0], 0]), rotation, hrz_transl_cum)

        # print('root', joint_coordinate['root'])
        
        idx = 0
        for idx_joint_angle, joint in enumerate(c_joints.values()):
            joint_angle = dec_Sxyz_ij[t, idx, :]
            child = joint
            if child.parent is not None:
                parent = child.parent
                # gets the coordinate in the current frame
                child_coordinate = - joint_angle*norm_ij[idx] + joint_coordinate[parent.name]
                joint_coordinate[child.name] = child_coordinate
                joint_coordinate_global[child.name] = inv_rotate_translate_frame(child_coordinate,  rotation, hrz_transl_cum)
                idx +=1

        joint_coordinate_arr.append(joint_coordinate_global)
    return joint_coordinate_arr


def video_generation(joint_coordinate_arr, c_joints, out_path = 'simple_animation.gif', display = "normal"):
    """
    display: normal, treadmill, circle
    """
    fig = plt.figure(figsize=(10, 10))
    # ax = Axes3D(fig)
    ax = fig.add_subplot(111, projection="3d")


    def draw_frame(i):
        ax.cla()
        ax.set_xlim3d(-50, 10)
        ax.set_ylim3d(-20, 40)
        ax.set_zlim3d(-20, 40)
        joint_coordinate = joint_coordinate_arr[i].copy()
        if display == "treadmill":
            offcet = joint_coordinate['root'].copy()
            offcet[1] = 0
            for joint_coord in joint_coordinate.values():
                joint_coord -= offcet
            # ax.plot(offcet[2], offcet[0], offcet[1], 'g.', markersize=10)
        
        if display == "circle":
            offcet = joint_coordinate['root'].copy()
            offcet[1] = 0
            # if it goes beyon the boundary it appears on the other side
            circle_coord = offcet.copy()
            
            circle_coord[2] = ((circle_coord[2] + 60) % 80) - 60
            circle_coord[0] = ((circle_coord[0] + 20) % 60) - 20

            for joint_coord in joint_coordinate.values():
                joint_coord -= offcet #+ circle_coord
                joint_coord += circle_coord
            # ax.plot(offcet[2], offcet[0], offcet[1], 'g.', markersize=10)
        
        # joints['root'].set_motion(motions[i])
        # c_joints = joints['root'].to_dict()
        xs, ys, zs = [], [], []
        for joint_coord in joint_coordinate.values():
            xs.append(joint_coord[0])
            ys.append(joint_coord[1])
            zs.append(joint_coord[2])
            ax.plot(zs, xs, ys, 'b.')

        for joint in c_joints.values():
            child = joint
            if child.parent is not None:
                parent = child.parent
                xs = [joint_coordinate[child.name][0], joint_coordinate[parent.name][0]]
                ys = [joint_coordinate[child.name][1], joint_coordinate[parent.name][1]]
                zs = [joint_coordinate[child.name][2], joint_coordinate[parent.name][2]]
                ax.plot(zs, xs, ys, 'r')
    FuncAnimation(fig, draw_frame, range(0, len(joint_coordinate_arr)-1, 10)).save(out_path, 
                                                    bitrate=8000,
                                                    fps=8)
    plt.close('all')
    FileLink(out_path)