motion_representation_amass.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from os.path import join as pjoin\n",
    "\n",
    "from common.skeleton import Skeleton\n",
    "import numpy as np\n",
    "import os\n",
    "from common.quaternion import *\n",
    "from paramUtil import *\n",
    "\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Apply the fix from [this issue](https://github.com/EricGuo5513/HumanML3D/issues/55) & [this issue](https://github.com/EricGuo5513/HumanML3D/issues/119) to Skeleton."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def uniform_skeleton(positions, target_offset):\n",
    "    src_skel = Skeleton(n_raw_offsets, kinematic_chain, 'cpu')\n",
    "    src_offset = src_skel.get_offsets_joints(torch.from_numpy(positions[0]))\n",
    "    src_offset = src_offset.numpy()\n",
    "    tgt_offset = target_offset.numpy()\n",
    "    # print(src_offset)\n",
    "    # print(tgt_offset)\n",
    "    '''Calculate Scale Ratio as the ratio of legs'''\n",
    "    src_leg_len = np.abs(src_offset[l_idx1]).max() + np.abs(src_offset[l_idx2]).max()\n",
    "    tgt_leg_len = np.abs(tgt_offset[l_idx1]).max() + np.abs(tgt_offset[l_idx2]).max()\n",
    "\n",
    "    scale_rt = tgt_leg_len / src_leg_len\n",
    "    # print(scale_rt)\n",
    "    src_root_pos = positions[:, 0]\n",
    "    tgt_root_pos = src_root_pos * scale_rt\n",
    "\n",
    "    '''Inverse Kinematics'''\n",
    "    quat_params = src_skel.inverse_kinematics_np(positions, face_joint_indx)\n",
    "    # print(quat_params.shape)\n",
    "\n",
    "    '''Forward Kinematics'''\n",
    "    src_skel.set_offset(target_offset)\n",
    "    new_joints = src_skel.forward_kinematics_np(quat_params, tgt_root_pos)\n",
    "    return new_joints\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_file(positions, feet_thre):\n",
    "    # (seq_len, joints_num, 3)\n",
    "    #     '''Down Sample'''\n",
    "    #     positions = positions[::ds_num]\n",
    "\n",
    "    '''Uniform Skeleton'''\n",
    "    positions = uniform_skeleton(positions, tgt_offsets)\n",
    "\n",
    "    '''Put on Floor'''\n",
    "    floor_height = positions.min(axis=0).min(axis=0)[1]\n",
    "    positions[:, :, 1] -= floor_height\n",
    "    #     print(floor_height)\n",
    "\n",
    "    #     plot_3d_motion(\"./positions_1.mp4\", kinematic_chain, positions, 'title', fps=20)\n",
    "\n",
    "    '''XZ at origin'''\n",
    "    root_pos_init = positions[0]\n",
    "    root_pose_init_xz = root_pos_init[0] * np.array([1, 0, 1])\n",
    "    positions = positions - root_pose_init_xz\n",
    "\n",
    "    # '''Move the first pose to origin '''\n",
    "    # root_pos_init = positions[0]\n",
    "    # positions = positions - root_pos_init[0]\n",
    "\n",
    "    '''All initially face Z+'''\n",
    "    r_hip, l_hip, sdr_r, sdr_l = face_joint_indx\n",
    "    across1 = root_pos_init[r_hip] - root_pos_init[l_hip]\n",
    "    across2 = root_pos_init[sdr_r] - root_pos_init[sdr_l]\n",
    "    across = across1 + across2\n",
    "    across = across / np.sqrt((across ** 2).sum(axis=-1))[..., np.newaxis]\n",
    "\n",
    "    # forward (3,), rotate around y-axis\n",
    "    forward_init = np.cross(np.array([[0, 1, 0]]), across, axis=-1)\n",
    "    # forward (3,)\n",
    "    forward_init = forward_init / np.sqrt((forward_init ** 2).sum(axis=-1))[..., np.newaxis]\n",
    "\n",
    "    #     print(forward_init)\n",
    "\n",
    "    target = np.array([[0, 0, 1]])\n",
    "    root_quat_init = qbetween_np(forward_init, target)\n",
    "    root_quat_init = np.ones(positions.shape[:-1] + (4,)) * root_quat_init\n",
    "\n",
    "    positions_b = positions.copy()\n",
    "\n",
    "    positions = qrot_np(root_quat_init, positions)\n",
    "\n",
    "    #     plot_3d_motion(\"./positions_2.mp4\", kinematic_chain, positions, 'title', fps=20)\n",
    "\n",
    "    '''New ground truth positions'''\n",
    "    global_positions = positions.copy()\n",
    "\n",
    "    # plt.plot(positions_b[:, 0, 0], positions_b[:, 0, 2], marker='*')\n",
    "    # plt.plot(positions[:, 0, 0], positions[:, 0, 2], marker='o', color='r')\n",
    "    # plt.xlabel('x')\n",
    "    # plt.ylabel('z')\n",
    "    # plt.axis('equal')\n",
    "    # plt.show()\n",
    "\n",
    "    \"\"\" Get Foot Contacts \"\"\"\n",
    "\n",
    "    def foot_detect(positions, thres):\n",
    "        velfactor, heightfactor = np.array([thres, thres]), np.array([3.0, 2.0])\n",
    "\n",
    "        feet_l_x = (positions[1:, fid_l, 0] - positions[:-1, fid_l, 0]) ** 2\n",
    "        feet_l_y = (positions[1:, fid_l, 1] - positions[:-1, fid_l, 1]) ** 2\n",
    "        feet_l_z = (positions[1:, fid_l, 2] - positions[:-1, fid_l, 2]) ** 2\n",
    "        #     feet_l_h = positions[:-1,fid_l,1]\n",
    "        #     feet_l = (((feet_l_x + feet_l_y + feet_l_z) < velfactor) & (feet_l_h < heightfactor)).astype(np.float)\n",
    "        feet_l = ((feet_l_x + feet_l_y + feet_l_z) < velfactor).astype(np.float32)\n",
    "\n",
    "        feet_r_x = (positions[1:, fid_r, 0] - positions[:-1, fid_r, 0]) ** 2\n",
    "        feet_r_y = (positions[1:, fid_r, 1] - positions[:-1, fid_r, 1]) ** 2\n",
    "        feet_r_z = (positions[1:, fid_r, 2] - positions[:-1, fid_r, 2]) ** 2\n",
    "        #     feet_r_h = positions[:-1,fid_r,1]\n",
    "        #     feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor) & (feet_r_h < heightfactor)).astype(np.float)\n",
    "        feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor)).astype(np.float32)\n",
    "        return feet_l, feet_r\n",
    "    #\n",
    "    feet_l, feet_r = foot_detect(positions, feet_thre)\n",
    "    # feet_l, feet_r = foot_detect(positions, 0.002)\n",
    "\n",
    "    '''Quaternion and Cartesian representation'''\n",
    "    r_rot = None\n",
    "\n",
    "    def get_rifke(positions):\n",
    "        '''Local pose'''\n",
    "        positions[..., 0] -= positions[:, 0:1, 0]\n",
    "        positions[..., 2] -= positions[:, 0:1, 2]\n",
    "        '''All pose face Z+'''\n",
    "        positions = qrot_np(np.repeat(r_rot[:, None], positions.shape[1], axis=1), positions)\n",
    "        return positions\n",
    "\n",
    "    def get_quaternion(positions):\n",
    "        skel = Skeleton(n_raw_offsets, kinematic_chain, \"cpu\")\n",
    "        # (seq_len, joints_num, 4)\n",
    "        quat_params = skel.inverse_kinematics_np(positions, face_joint_indx, smooth_forward=False)\n",
    "\n",
    "        '''Fix Quaternion Discontinuity'''\n",
    "        quat_params = qfix(quat_params)\n",
    "        # (seq_len, 4)\n",
    "        r_rot = quat_params[:, 0].copy()\n",
    "        #     print(r_rot[0])\n",
    "        '''Root Linear Velocity'''\n",
    "        # (seq_len - 1, 3)\n",
    "        velocity = (positions[1:, 0] - positions[:-1, 0]).copy()\n",
    "        #     print(r_rot.shape, velocity.shape)\n",
    "        velocity = qrot_np(r_rot[1:], velocity)\n",
    "        '''Root Angular Velocity'''\n",
    "        # (seq_len - 1, 4)\n",
    "        r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1]))\n",
    "        quat_params[1:, 0] = r_velocity\n",
    "        # (seq_len, joints_num, 4)\n",
    "        return quat_params, r_velocity, velocity, r_rot\n",
    "\n",
    "    def get_cont6d_params(positions):\n",
    "        skel = Skeleton(n_raw_offsets, kinematic_chain, \"cpu\")\n",
    "        # (seq_len, joints_num, 4)\n",
    "        quat_params = skel.inverse_kinematics_np(positions, face_joint_indx, smooth_forward=True)\n",
    "\n",
    "        '''Quaternion to continuous 6D'''\n",
    "        cont_6d_params = quaternion_to_cont6d_np(quat_params)\n",
    "        # (seq_len, 4)\n",
    "        r_rot = quat_params[:, 0].copy()\n",
    "        #     print(r_rot[0])\n",
    "        '''Root Linear Velocity'''\n",
    "        # (seq_len - 1, 3)\n",
    "        velocity = (positions[1:, 0] - positions[:-1, 0]).copy()\n",
    "        #     print(r_rot.shape, velocity.shape)\n",
    "        velocity = qrot_np(r_rot[1:], velocity)\n",
    "        '''Root Angular Velocity'''\n",
    "        # (seq_len - 1, 4)\n",
    "        r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1]))\n",
    "        # (seq_len, joints_num, 4)\n",
    "        return cont_6d_params, r_velocity, velocity, r_rot\n",
    "\n",
    "    cont_6d_params, r_velocity, velocity, r_rot = get_cont6d_params(positions)\n",
    "    positions = get_rifke(positions)\n",
    "\n",
    "    #     trejec = np.cumsum(np.concatenate([np.array([[0, 0, 0]]), velocity], axis=0), axis=0)\n",
    "    #     r_rotations, r_pos = recover_ric_glo_np(r_velocity, velocity[:, [0, 2]])\n",
    "\n",
    "    # plt.plot(positions_b[:, 0, 0], positions_b[:, 0, 2], marker='*')\n",
    "    # plt.plot(ground_positions[:, 0, 0], ground_positions[:, 0, 2], marker='o', color='r')\n",
    "    # plt.plot(trejec[:, 0], trejec[:, 2], marker='^', color='g')\n",
    "    # plt.plot(r_pos[:, 0], r_pos[:, 2], marker='s', color='y')\n",
    "    # plt.xlabel('x')\n",
    "    # plt.ylabel('z')\n",
    "    # plt.axis('equal')\n",
    "    # plt.show()\n",
    "\n",
    "    '''Root height'''\n",
    "    root_y = positions[:, 0, 1:2]\n",
    "\n",
    "    '''Root rotation and linear velocity'''\n",
    "    # (seq_len-1, 1) rotation velocity along y-axis\n",
    "    # (seq_len-1, 2) linear velovity on xz plane\n",
    "    r_velocity = np.arcsin(r_velocity[:, 2:3])\n",
    "    l_velocity = velocity[:, [0, 2]]\n",
    "    #     print(r_velocity.shape, l_velocity.shape, root_y.shape)\n",
    "    root_data = np.concatenate([r_velocity, l_velocity, root_y[:-1]], axis=-1)\n",
    "\n",
    "    '''Get Joint Rotation Representation'''\n",
    "    # (seq_len, (joints_num-1) *6) quaternion for skeleton joints\n",
    "    rot_data = cont_6d_params[:, 1:].reshape(len(cont_6d_params), -1)\n",
    "\n",
    "    '''Get Joint Rotation Invariant Position Represention'''\n",
    "    # (seq_len, (joints_num-1)*3) local joint position\n",
    "    ric_data = positions[:, 1:].reshape(len(positions), -1)\n",
    "\n",
    "    '''Get Joint Velocity Representation'''\n",
    "    # (seq_len-1, joints_num*3)\n",
    "    local_vel = qrot_np(np.repeat(r_rot[:-1, None], global_positions.shape[1], axis=1),\n",
    "                        global_positions[1:] - global_positions[:-1])\n",
    "    local_vel = local_vel.reshape(len(local_vel), -1)\n",
    "\n",
    "    data = root_data\n",
    "    data = np.concatenate([data, ric_data[:-1]], axis=-1)\n",
    "    data = np.concatenate([data, rot_data[:-1]], axis=-1)\n",
    "    #     print(data.shape, local_vel.shape)\n",
    "    data = np.concatenate([data, local_vel], axis=-1)\n",
    "    data = np.concatenate([data, feet_l, feet_r], axis=-1)\n",
    "\n",
    "    return data, global_positions, positions, l_velocity\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Recover global angle and positions for rotation data\n",
    "# root_rot_velocity (B, seq_len, 1)\n",
    "# root_linear_velocity (B, seq_len, 2)\n",
    "# root_y (B, seq_len, 1)\n",
    "# ric_data (B, seq_len, (joint_num - 1)*3)\n",
    "# rot_data (B, seq_len, (joint_num - 1)*6)\n",
    "# local_velocity (B, seq_len, joint_num*3)\n",
    "# foot contact (B, seq_len, 4)\n",
    "def recover_root_rot_pos(data):\n",
    "    rot_vel = data[..., 0]\n",
    "    r_rot_ang = torch.zeros_like(rot_vel).to(data.device)\n",
    "    '''Get Y-axis rotation from rotation velocity'''\n",
    "    r_rot_ang[..., 1:] = rot_vel[..., :-1]\n",
    "    r_rot_ang = torch.cumsum(r_rot_ang, dim=-1)\n",
    "\n",
    "    r_rot_quat = torch.zeros(data.shape[:-1] + (4,)).to(data.device)\n",
    "    r_rot_quat[..., 0] = torch.cos(r_rot_ang)\n",
    "    r_rot_quat[..., 2] = torch.sin(r_rot_ang)\n",
    "\n",
    "    r_pos = torch.zeros(data.shape[:-1] + (3,)).to(data.device)\n",
    "    r_pos[..., 1:, [0, 2]] = data[..., :-1, 1:3]\n",
    "    '''Add Y-axis rotation to root position'''\n",
    "    r_pos = qrot(qinv(r_rot_quat), r_pos)\n",
    "\n",
    "    r_pos = torch.cumsum(r_pos, dim=-2)\n",
    "\n",
    "    r_pos[..., 1] = data[..., 3]\n",
    "    return r_rot_quat, r_pos\n",
    "\n",
    "\n",
    "def recover_from_rot(data, joints_num, skeleton):\n",
    "    r_rot_quat, r_pos = recover_root_rot_pos(data)\n",
    "\n",
    "    r_rot_cont6d = quaternion_to_cont6d(r_rot_quat)\n",
    "\n",
    "    start_indx = 1 + 2 + 1 + (joints_num - 1) * 3\n",
    "    end_indx = start_indx + (joints_num - 1) * 6\n",
    "    cont6d_params = data[..., start_indx:end_indx]\n",
    "    #     print(r_rot_cont6d.shape, cont6d_params.shape, r_pos.shape)\n",
    "    cont6d_params = torch.cat([r_rot_cont6d, cont6d_params], dim=-1)\n",
    "    cont6d_params = cont6d_params.view(-1, joints_num, 6)\n",
    "\n",
    "    positions = skeleton.forward_kinematics_cont6d(cont6d_params, r_pos)\n",
    "\n",
    "    return positions\n",
    "\n",
    "\n",
    "def recover_from_ric(data, joints_num):\n",
    "    r_rot_quat, r_pos = recover_root_rot_pos(data)\n",
    "    positions = data[..., 4:(joints_num - 1) * 3 + 4]\n",
    "    positions = positions.view(positions.shape[:-1] + (-1, 3))\n",
    "\n",
    "    '''Add Y-axis rotation to local joints'''\n",
    "    positions = qrot(qinv(r_rot_quat[..., None, :]).expand(positions.shape[:-1] + (4,)), positions)\n",
    "\n",
    "    '''Add root XZ to joints'''\n",
    "    positions[..., 0] += r_pos[..., 0:1]\n",
    "    positions[..., 2] += r_pos[..., 2:3]\n",
    "\n",
    "    '''Concate root and joints'''\n",
    "    positions = torch.cat([r_pos.unsqueeze(-2), positions], dim=-2)\n",
    "\n",
    "    return positions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The given data is used to double check if you are on the right track.\n",
    "reference1 = np.load('./HumanML3D/new_joints/012314.npy')\n",
    "reference2 = np.load('./HumanML3D/new_joint_vecs/012314.npy')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `example_id` below is randomly selected. Explanation in [this issue](https://github.com/EricGuo5513/HumanML3D/issues/37)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                                                                                       | 0/16118 [00:00<?, ?it/s]/vision/changan/shrinidhi/HumanML3D/common/quaternion.py:398: UserWarning: Using torch.cross without specifying the dim arg is deprecated.\n",
      "Please either pass the dim explicitly or simply use torch.linalg.cross.\n",
      "The default value of dim will change to agree with that of linalg.cross in a future release. (Triggered internally at ../aten/src/ATen/native/Cross.cpp:62.)\n",
      "  v = torch.cross(v0, v1)\n",
      "  4%|████▏                                                                                                        | 627/16118 [00:37<10:52, 23.73it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "000281_02.npy\n",
      "index 0 is out of bounds for axis 0 with size 0\n",
      "M000281_02.npy\n",
      "index 0 is out of bounds for axis 0 with size 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 19%|████████████████████▏                                                                                       | 3011/16118 [02:49<09:16, 23.57it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "001497_02.npy\n",
      "index 0 is out of bounds for axis 0 with size 0\n",
      "M001497_02.npy\n",
      "index 0 is out of bounds for axis 0 with size 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|█████████████████████▏                                                                                      | 3161/16118 [02:57<09:28, 22.78it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "001570_02.npy\n",
      "index 0 is out of bounds for axis 0 with size 0\n",
      "M001570_02.npy\n",
      "index 0 is out of bounds for axis 0 with size 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 25%|███████████████████████████▏                                                                                | 4063/16118 [03:46<07:06, 28.27it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "002032_02.npy\n",
      "index 0 is out of bounds for axis 0 with size 0\n",
      "M002032_02.npy\n",
      "index 0 is out of bounds for axis 0 with size 0\n",
      "002032_03.npy\n",
      "index 0 is out of bounds for axis 0 with size 0\n",
      "M002032_03.npy\n",
      "index 0 is out of bounds for axis 0 with size 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 91%|█████████████████████████████████████████████████████████████████████████████████████████████████▏         | 14647/16118 [14:06<01:11, 20.54it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "009471_00.npy\n",
      "index 0 is out of bounds for axis 0 with size 0\n",
      "M009471_00.npy\n",
      "index 0 is out of bounds for axis 0 with size 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 16118/16118 [15:32<00:00, 17.28it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total clips: 16118, Frames: 2463964, Duration: 2053.303333m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "'''\n",
    "For HumanML3D Dataset\n",
    "'''\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    example_id = \"000021_00\"\n",
    "    # Lower legs\n",
    "    l_idx1, l_idx2 = 5, 8\n",
    "    # Right/Left foot\n",
    "    fid_r, fid_l = [8, 11], [7, 10]\n",
    "    # Face direction, r_hip, l_hip, sdr_r, sdr_l\n",
    "    face_joint_indx = [2, 1, 17, 16]\n",
    "    # l_hip, r_hip\n",
    "    r_hip, l_hip = 2, 1\n",
    "    joints_num = 22\n",
    "    # ds_num = 8\n",
    "    data_dir = './joints_amass/'\n",
    "    save_dir1 = './HumanML3D_amass/new_joints/'\n",
    "    save_dir2 = './HumanML3D_amass/new_joint_vecs/'\n",
    "    \n",
    "    os.makedirs(save_dir1, exist_ok=True)\n",
    "    os.makedirs(save_dir2, exist_ok=True)\n",
    "\n",
    "    n_raw_offsets = torch.from_numpy(t2m_raw_offsets)\n",
    "    kinematic_chain = t2m_kinematic_chain\n",
    "\n",
    "    # Get offsets of target skeleton\n",
    "    example_data = np.load(os.path.join(data_dir, example_id + '.npy'))\n",
    "    example_data = example_data.reshape(len(example_data), -1, 3)\n",
    "    example_data = torch.from_numpy(example_data)\n",
    "    tgt_skel = Skeleton(n_raw_offsets, kinematic_chain, 'cpu')\n",
    "    # (joints_num, 3)\n",
    "    tgt_offsets = tgt_skel.get_offsets_joints(example_data[0])\n",
    "    # print(tgt_offsets)\n",
    "\n",
    "    source_list = os.listdir(data_dir)\n",
    "    frame_num = 0\n",
    "    for source_file in tqdm(source_list,ncols=150):\n",
    "        source_data = np.load(os.path.join(data_dir, source_file))[:, :joints_num]\n",
    "        try:\n",
    "            data, ground_positions, positions, l_velocity = process_file(source_data, 0.002)\n",
    "            rec_ric_data = recover_from_ric(torch.from_numpy(data).unsqueeze(0).float(), joints_num)\n",
    "            np.save(pjoin(save_dir1, source_file), rec_ric_data.squeeze().numpy())\n",
    "            np.save(pjoin(save_dir2, source_file), data)\n",
    "            frame_num += data.shape[0]\n",
    "        except Exception as e:\n",
    "            print(source_file)\n",
    "            print(e)\n",
    "#         print(source_file)\n",
    "#         break\n",
    "\n",
    "    print('Total clips: %d, Frames: %d, Duration: %fm' %\n",
    "          (len(source_list), frame_num, frame_num / 20 / 60))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Drop the files with NaN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pathlib\n",
    "\n",
    "new_joint_vecs_d = pathlib.Path('HumanML3D_amass/new_joint_vecs')\n",
    "n_dropped = 0\n",
    "for f in tqdm(list(new_joint_vecs_d.iterdir()),ncols=150):\n",
    "    if np.isnan(np.load(f)).any():\n",
    "        f.unlink()\n",
    "        n_dropped += 1\n",
    "print('files removed:',n_dropped)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Drop the ids from train/test set if there is no corresponding file. Also there are duplicate ids in the test set of AMASS. De-duplicate it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_f = pathlib.Path('HumanML3D_amass/train.txt')\n",
    "test_f = pathlib.Path('HumanML3D_amass/test.txt')\n",
    "train = [id.strip()\n",
    "         for id in open(train_f)]\n",
    "new_train = {id\n",
    "             for id in train\n",
    "             if (new_joint_vecs_d / f'{id}.npy').is_file()}\n",
    "print('train:',len(train))\n",
    "print('new train:',len(new_train))\n",
    "\n",
    "test = [id.strip()\n",
    "        for id in open(test_f)]\n",
    "new_test = {id\n",
    "            for id in test\n",
    "            if (new_joint_vecs_d / f'{id}.npy').is_file()}\n",
    "print('test:',len(test))\n",
    "print('new test:',len(new_test))\n",
    "\n",
    "_ = open(train_f,'w').write('\\n'.join(sorted(new_train))+'\\n')\n",
    "_ = open(test_f,'w').write('\\n'.join(sorted(new_test))+'\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check if your data is correct. If it's aligned with the given reference, then it is right"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reference1_1 = np.load('./HumanML3D/new_joints/012314.npy')\n",
    "reference2_1 = np.load('./HumanML3D/new_joint_vecs/012314.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "abs(reference1 - reference1_1).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "abs(reference2 - reference2_1).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}