{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbfb8da4-14a5-44d8-b9b0-d66f032f09fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.nice(5)\n",
    "import rerun as rr\n",
    "import numpy as np\n",
    "from math import tan\n",
    "import time\n",
    "from utils import remake_dir\n",
    "import pandas as pd\n",
    "from data_loaders.ManoHandDataProvider import MANOHandDataProvider\n",
    "from data_loaders.loader_object_library import load_object_library\n",
    "from data_loaders.mano_layer import MANOHandModel\n",
    "from data_loaders.loader_hand_poses import Handedness, HandPose\n",
    "from data_loaders.hand_common import LANDMARK_CONNECTIVITY\n",
    "from data_loaders.loader_object_library import ObjectLibrary\n",
    "from data_loaders.headsets import Headset\n",
    "from projectaria_tools.core.stream_id import StreamId\n",
    "from projectaria_tools.core.sensor_data import TimeDomain, TimeQueryOptions\n",
    "from projectaria_tools.core.sophus import SE3\n",
    "from projectaria_tools.utils.rerun_helpers import ToTransform3D\n",
    "\n",
    "\n",
    "data_path = '/scratch/hu/pose_forecast/hot3d_hoigaze/P0001_10a27bf7_room_721_890_'\n",
    "timestamps_path = data_path + 'timestamps.npy'\n",
    "head_path = data_path + 'head.npy'\n",
    "gaze_path = data_path + 'gaze.npy'\n",
    "hand_path = data_path + 'hand.npy'\n",
    "hand_joint_path = data_path + 'handjoints.npy'\n",
    "object_path = data_path + 'objects.npy'\n",
    "object_bbx_path = data_path + 'object_bbx.npy'\n",
    "object_library_path = '/datasets/public/zhiming_datasets/hot3d/assets/'\n",
    "mano_hand_model_path = '/datasets/public/zhiming_datasets/hot3d/mano_v1_2/models/'\n",
    "\n",
    "show_bbx = False\n",
    "show_hand_mesh = True\n",
    "# init the object library\n",
    "if not os.path.exists(object_library_path):\n",
    "    print(\"invalid object library path.\")\n",
    "    print(\"please follow the instructions at https://github.com/facebookresearch/hot3d to Download the HOT3D Assets Dataset\")        \n",
    "    raise    \n",
    "object_library = load_object_library(object_library_folderpath=object_library_path)\n",
    "\n",
    "# init the HANDs model\n",
    "if not os.path.exists(mano_hand_model_path):\n",
    "    print(\"invalid mano hand model path.\")\n",
    "    print(\"please follow the instructions at https://github.com/facebookresearch/hot3d to Download the MANO files\")\n",
    "    raise\n",
    "mano_hand_model = MANOHandModel(mano_hand_model_path)\n",
    "\n",
    "timestamps_data = np.load(timestamps_path)\n",
    "head_data = np.load(head_path) # head_direction (3) + head_translation (3) + head_rotation (4, quat_xyzw)\n",
    "gaze_data = np.load(gaze_path) # gaze_direction (3) + gaze_center_in_world (3)\n",
    "hand_data = np.load(hand_path) # left_hand (22) + right_hand (22), hand = wrist_pose (7, translation (3) + rotation (4)) + joint_angles (15)\n",
    "hand_joint_data = np.load(hand_joint_path) # left_hand (20*3) + right_hand (20*3)\n",
    "object_data = np.load(object_path) # object_data (8) * 6 objects (at most 6 objects), object_data = object_uid (1) + object_pose (7, translation (3) + rotation (4))\n",
    "object_bbx_data = np.load(object_bbx_path) # bounding box information: 6 objects (at most 6 objects) * 8 vertexes * 3\n",
    "frame_length = len(head_data)\n",
    "\n",
    "# alias over the HAND pose data provider\n",
    "hand_data_provider = MANOHandDataProvider('./mano_hand_pose_init/mano_hand_pose_trajectory.jsonl', mano_hand_model)\n",
    "# keep track of what 3D assets have been loaded/unloaded so we will load them only when needed\n",
    "object_cache_status = {}\n",
    "\n",
    "# Init a rerun context\n",
    "rr.init(\"hot3d-aria\")\n",
    "rec = rr.memory_recording()\n",
    "\n",
    "def log_pose(\n",
    "    pose: SE3,\n",
    "    label: str,\n",
    "    static=False\n",
    ") -> None:\n",
    "    rr.log(label, ToTransform3D(pose, False), static=static)\n",
    "\n",
    "\n",
    "for i in range(frame_length):\n",
    "    timestamp_ns = timestamps_data[i, 0]\n",
    "    rr.set_time_nanos(\"synchronization_time\", int(timestamp_ns))\n",
    "    rr.set_time_sequence(\"timestamp\", int(timestamp_ns))\n",
    "    \n",
    "    head_direction = head_data[i, 0:3]\n",
    "    head_translation = head_data[i, 3:6]\n",
    "    head_rotation = head_data[i, 6:10]    \n",
    "    gaze_direction = gaze_data[i, 0:3]\n",
    "    gaze_center_in_world = gaze_data[i, 3:6]\n",
    "    left_hand_translation = hand_data[i, 0:3]\n",
    "    left_hand_rotation = hand_data[i, 3:7]\n",
    "    left_hand_joint_angles = hand_data[i, 7:22]\n",
    "    left_hand_joints = hand_joint_data[i, 0:60].reshape((20, 3))\n",
    "    right_hand_translation = hand_data[i, 22:25]\n",
    "    right_hand_rotation = hand_data[i, 25:29]\n",
    "    right_hand_joint_angles = hand_data[i, 29:44]\n",
    "    right_hand_joints = hand_joint_data[i, 60:120].reshape((20, 3))\n",
    "    left_hand_wrist_pose = SE3.from_quat_and_translation(left_hand_rotation[-1], left_hand_rotation[:-1], left_hand_translation)\n",
    "    left_hand_pose = HandPose(Handedness.Left, left_hand_wrist_pose, left_hand_joint_angles)\n",
    "    right_hand_wrist_pose = SE3.from_quat_and_translation(right_hand_rotation[-1], right_hand_rotation[:-1], right_hand_translation)\n",
    "    right_hand_pose = HandPose(Handedness.Right, right_hand_wrist_pose, right_hand_joint_angles)\n",
    "    \n",
    "    # use cpf pose as head pose, see https://facebookresearch.github.io/projectaria_tools/docs/data_formats/coordinate_convention/3d_coordinate_frame_convention\n",
    "    T_world_cpf = SE3.from_quat_and_translation(head_rotation[-1], head_rotation[:-1], head_translation)\n",
    "    log_pose(pose=T_world_cpf, label=\"world/head_pose\")\n",
    "    #rr.log(\n",
    "    #\"world/head_direction\",\n",
    "    #rr.Points3D([head_translation], radii=[0.003]),\n",
    "    #rr.Arrows3D(vectors=[head_direction*0.4], colors=[[0, 0.8, 0.8, 0.5]]))\n",
    "    #log_pose(pose=left_hand_wrist_pose, label=\"world/left_hand_pose\")\n",
    "    #log_pose(pose=right_hand_wrist_pose, label=\"world/right_hand_pose\")\n",
    "    \n",
    "    rr.log(\n",
    "    \"world/gaze_direction\",\n",
    "    rr.Points3D([head_translation], radii=[0.003]),\n",
    "    rr.Arrows3D(vectors=[gaze_direction*0.4], colors=[[0, 0.8, 0.2, 0.5]]))\n",
    "    #print(\"frame: {}, gaze: {}\".format(i+1119, gaze_direction))\n",
    "    \n",
    "    # plot hands as a triangular mesh representation\n",
    "    if show_hand_mesh:\n",
    "        left_hand_mesh_vertices = hand_data_provider.get_hand_mesh_vertices(left_hand_pose)\n",
    "        left_hand_triangles, left_hand_vertex_normals = hand_data_provider.get_hand_mesh_faces_and_normals(left_hand_pose)    \n",
    "        rr.log(\n",
    "            f\"world/left_hand/mesh_faces\",\n",
    "            rr.Mesh3D(\n",
    "                vertex_positions=left_hand_mesh_vertices,\n",
    "                vertex_normals=left_hand_vertex_normals,\n",
    "                triangle_indices=left_hand_triangles))\n",
    "        right_hand_mesh_vertices = hand_data_provider.get_hand_mesh_vertices(right_hand_pose)\n",
    "        right_hand_triangles, right_hand_vertex_normals = hand_data_provider.get_hand_mesh_faces_and_normals(right_hand_pose)\n",
    "        rr.log(\n",
    "            f\"world/right_hand/mesh_faces\",\n",
    "            rr.Mesh3D(\n",
    "                vertex_positions=right_hand_mesh_vertices,\n",
    "                vertex_normals=right_hand_vertex_normals,\n",
    "                triangle_indices=right_hand_triangles))\n",
    "    else:\n",
    "        #left_hand_translation = np.array([0, 0, 0])\n",
    "        #left_hand_rotation = np.array([0, 0, 0, 1])\n",
    "        #left_hand_joint_angles = np.zeros(15)\n",
    "        #left_hand_wrist_pose = SE3.from_quat_and_translation(left_hand_rotation[-1], left_hand_rotation[:-1], left_hand_translation)\n",
    "        #log_pose(pose=left_hand_wrist_pose, label=\"world/left_hand_pose\")\n",
    "        #left_hand_pose = HandPose(Handedness.Left, left_hand_wrist_pose, left_hand_joint_angles)\n",
    "        #left_hand_joints = hand_data_provider.get_hand_landmarks(left_hand_pose)\n",
    "        #left_hand_wrist = left_hand_joints[5, :].clone()\n",
    "        #joint_number = left_hand_joints.shape[0]\n",
    "        #for index in range(joint_number):\n",
    "        #    left_hand_joints[index, :] -= left_hand_wrist\n",
    "        #for index in range(joint_number):            \n",
    "        #    tmp = left_hand_joints[index, :].clone()\n",
    "        #    left_hand_joints[index, 1] = -tmp[2]\n",
    "        #    left_hand_joints[index, 2] = tmp[1]\n",
    "        #for index in range(joint_number):            \n",
    "        #    print(left_hand_joints[index])\n",
    "        \n",
    "        #right_hand_translation = np.array([0, 0, 0])\n",
    "        #right_hand_rotation = np.array([0, 0, 0, 1])\n",
    "        #right_hand_joint_angles = np.zeros(15)\n",
    "        #right_hand_wrist_pose = SE3.from_quat_and_translation(right_hand_rotation[-1], right_hand_rotation[:-1], right_hand_translation)\n",
    "        #log_pose(pose=right_hand_wrist_pose, label=\"world/right_hand_pose\")\n",
    "        #right_hand_pose = HandPose(Handedness.Right, right_hand_wrist_pose, right_hand_joint_angles)\n",
    "        #right_hand_joints = hand_data_provider.get_hand_landmarks(right_hand_pose)\n",
    "        #right_hand_wrist = right_hand_joints[5, :].clone()\n",
    "        #joint_number = right_hand_joints.shape[0]\n",
    "        #for index in range(joint_number):\n",
    "        #    right_hand_joints[index, :] -= right_hand_wrist\n",
    "        #for index in range(joint_number):            \n",
    "        #    tmp = right_hand_joints[index, :].clone()\n",
    "        #    right_hand_joints[index, 1] = -tmp[2]\n",
    "        #    right_hand_joints[index, 2] = tmp[1]\n",
    "        #for index in range(joint_number):            \n",
    "        #    print(right_hand_joints[index])\n",
    "        \n",
    "        left_hand_skeleton = [connections\n",
    "                  for connectivity in LANDMARK_CONNECTIVITY\n",
    "                  for connections in [[left_hand_joints[it].tolist() for it in connectivity]]]                \n",
    "        rr.log(\n",
    "            f\"world/left_hand_skeleton\",\n",
    "            rr.LineStrips3D(left_hand_skeleton, radii=0.002),\n",
    "        )\n",
    "        right_hand_skeleton = [connections\n",
    "                  for connectivity in LANDMARK_CONNECTIVITY\n",
    "                  for connections in [[right_hand_joints[it].tolist() for it in connectivity]]]        \n",
    "        rr.log(\n",
    "            f\"world/right_hand_skeleton\",\n",
    "            rr.LineStrips3D(right_hand_skeleton, radii=0.002),\n",
    "        )\n",
    "        \n",
    "    # load objects\n",
    "    object_num_max = 6\n",
    "    logging_status = {}    \n",
    "    for item in range(object_num_max):\n",
    "        object_uid = str(int(object_data[i, item*8]))\n",
    "        if object_uid == '0':\n",
    "            break\n",
    "        logging_status[object_uid] = False\n",
    "    object_num = len(logging_status)\n",
    "    for item in range(object_num):\n",
    "        object_uid = str(int(object_data[i, item*8]))\n",
    "        object_name = object_library.object_id_to_name_dict[object_uid]\n",
    "        object_cad_asset_filepath = ObjectLibrary.get_cad_asset_path(\n",
    "            object_library_folderpath=object_library.asset_folder_name,\n",
    "            object_id=object_uid)\n",
    "        object_translation = object_data[i, item*8+1:item*8+4]\n",
    "        object_rotation = object_data[i, item*8+4:item*8+8]\n",
    "        object_pose = SE3.from_quat_and_translation(object_rotation[-1], object_rotation[:-1], object_translation)        \n",
    "        log_pose(pose=object_pose, label=f\"world/objects/{object_name}\")\n",
    "        logging_status[object_uid] = True # mark object has been seen (enable to know which object has been logged or not)\n",
    "        \n",
    "        if show_bbx:\n",
    "            bbx_vertex_0 = object_bbx_data[i, item*24:item*24+3]\n",
    "            bbx_vertex_1 = object_bbx_data[i, item*24+3:item*24+6]\n",
    "            bbx_vertex_2 = object_bbx_data[i, item*24+6:item*24+9]\n",
    "            bbx_vertex_3 = object_bbx_data[i, item*24+9:item*24+12]\n",
    "            bbx_vertex_4 = object_bbx_data[i, item*24+12:item*24+15]\n",
    "            bbx_vertex_5 = object_bbx_data[i, item*24+15:item*24+18]\n",
    "            bbx_vertex_6 = object_bbx_data[i, item*24+18:item*24+21]\n",
    "            bbx_vertex_7 = object_bbx_data[i, item*24+21:item*24+24]    \n",
    "            points = [\n",
    "                bbx_vertex_0,\n",
    "                bbx_vertex_1,\n",
    "                bbx_vertex_2,\n",
    "                bbx_vertex_3,\n",
    "                bbx_vertex_0,\n",
    "                bbx_vertex_7,\n",
    "                bbx_vertex_6,\n",
    "                bbx_vertex_5,\n",
    "                bbx_vertex_4,\n",
    "                bbx_vertex_7,\n",
    "                bbx_vertex_6,\n",
    "                bbx_vertex_1,\n",
    "                bbx_vertex_2,\n",
    "                bbx_vertex_5,\n",
    "                bbx_vertex_4,\n",
    "                bbx_vertex_3]\n",
    "            rr.log(f\"world/objects_bbx/{object_name}\", rr.LineStrips3D([points]))\n",
    "            \n",
    "        # link the corresponding 3D object to the pose\n",
    "        if object_uid not in object_cache_status.keys():\n",
    "            object_cache_status[object_uid] = True            \n",
    "            rr.log(\n",
    "                f\"world/objects/{object_name}\",\n",
    "                rr.Asset3D(path=object_cad_asset_filepath))            \n",
    "    # if some objects are not visible, we clear the entity\n",
    "    for object_uid, displayed in logging_status.items():\n",
    "        if not displayed:\n",
    "            object_name = object_library.object_id_to_name_dict[object_uid]\n",
    "            rr.log(\n",
    "                f\"world/objects/{object_name}\",\n",
    "                rr.Clear.recursive())\n",
    "            if show_bbx:\n",
    "                rr.log(\n",
    "                    f\"world/objects_bbx/{object_name}\",\n",
    "                    rr.Clear.recursive())            \n",
    "            if object_uid in object_cache_status.keys():\n",
    "                del object_cache_status[object_uid]  # we will log the mesh again\n",
    "                \n",
    "# show the rerun window\n",
    "rr.notebook_show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7811d038-6150-46b3-ac62-d4875191e0f4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
