{
 "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 dataset_api import Hot3dDataProvider\n",
    "from data_loaders.loader_object_library import load_object_library\n",
    "from data_loaders.mano_layer import MANOHandModel\n",
    "from data_loaders.loader_masks import combine_mask_data, load_mask_data, MaskData\n",
    "from data_loaders.loader_hand_poses import Handedness, HandPose\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",
    "\n",
    "\n",
    "dataset_path = '/datasets/public/zhiming_datasets/hot3d/aria/'\n",
    "dataset_processed_path = '/scratch/hu/pose_forecast/hot3d_hoigaze/'\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",
    "remake_dir(dataset_processed_path)\n",
    "dataset_info = pd.read_csv('hot3d_aria_scene.csv')\n",
    "valid_frame_length = 60 # 60 frames -> 2 seconds, dropout the recordings that are too short\n",
    "\n",
    "# init the object library\n",
    "if not os.path.exists(object_library_path):\n",
    "    print(\"invalid library path.\")\n",
    "    print(\"please do update the path to VALID values for your system.\")\n",
    "    raise\n",
    "object_library = load_object_library(object_library_folderpath=object_library_path)\n",
    "\n",
    "# load the bounding box information of the objects\n",
    "object_info = pd.read_csv('hot3d_objects.csv')\n",
    "object_bbx = {}\n",
    "for i, uid in enumerate(object_info['object_uid']):    \n",
    "    bbx_x_min = object_info['bbx_x_min'][i]\n",
    "    bbx_x_max = object_info['bbx_x_max'][i]\n",
    "    bbx_y_min = object_info['bbx_y_min'][i]\n",
    "    bbx_y_max = object_info['bbx_y_max'][i]\n",
    "    bbx_z_min = object_info['bbx_z_min'][i]\n",
    "    bbx_z_max = object_info['bbx_z_max'][i]\n",
    "    bbx = [bbx_x_min, bbx_x_max, bbx_y_min, bbx_y_max, bbx_z_min, bbx_z_max]\n",
    "    object_bbx[str(uid)] = bbx\n",
    "    \n",
    "# init the HANDs model. If None, the UmeTrack HANDs model will be used\n",
    "mano_hand_model = None\n",
    "if mano_hand_model_path is not None:\n",
    "    mano_hand_model = MANOHandModel(mano_hand_model_path)\n",
    "    \n",
    "for i, seq in enumerate(dataset_info['sequence_name']):\n",
    "    scene = dataset_info['scene'][i]\n",
    "    print(\"\\nprocessing {}th seq: {}, scene: {}...\".format(i+1, seq, scene))\n",
    "    seq_path = dataset_path + seq + '/'\n",
    "    if not os.path.exists(seq_path):\n",
    "        print(\"invalid input sequence path.\")\n",
    "        print(\"please do update the path to VALID values for your system.\")\n",
    "        raise\n",
    "    save_path = dataset_processed_path + seq + '_' + scene + '_'\n",
    "\n",
    "    # segment the sequence into valid and invalid parts using the masks    \n",
    "    mask_list = [\n",
    "        \"masks/mask_object_pose_available.csv\",\n",
    "        \"masks/mask_hand_pose_available.csv\",    \n",
    "        \"masks/mask_headset_pose_available.csv\",\n",
    "        #\"masks/mask_object_visibility.csv\",    \n",
    "        #\"masks/mask_hand_visible.csv\",\n",
    "        #\"masks/mask_good_exposure.csv\",\n",
    "        \"masks/mask_qa_pass.csv\"]\n",
    "    \n",
    "    # load the referred masks\n",
    "    mask_data_list = []\n",
    "    for it in mask_list:\n",
    "        if os.path.exists(os.path.join(seq_path, it)):\n",
    "            ret = load_mask_data(os.path.join(seq_path, it))\n",
    "            mask_data_list.append(ret)\n",
    "    # combine the masks (you can choose logical \"and\"/\"or\")\n",
    "    combined_masks = combine_mask_data(mask_data_list, \"and\")\n",
    "    masks = []\n",
    "    for value in combined_masks.data['214-1'].values():\n",
    "        masks.append(value)\n",
    "    print(\"valid frames: {}/{}\".format(sum(masks), len(masks)))\n",
    "    \n",
    "    # initialize hot3d data provider\n",
    "    hot3d_data_provider = Hot3dDataProvider(\n",
    "    sequence_folder=seq_path,\n",
    "    object_library=object_library,\n",
    "    mano_hand_model=mano_hand_model)\n",
    "    #print(f\"data_provider statistics: {hot3d_data_provider.get_data_statistics()}\")        \n",
    "    # alias over the HAND pose data provider\n",
    "    hand_data_provider = hot3d_data_provider.mano_hand_data_provider if hot3d_data_provider.mano_hand_data_provider is not None else hot3d_data_provider.umetrack_hand_data_provider\n",
    "    # alias over the Object pose data provider\n",
    "    object_pose_data_provider = hot3d_data_provider.object_pose_data_provider\n",
    "    # alias over the HEADSET/Device pose data provider\n",
    "    device_pose_provider = hot3d_data_provider.device_pose_data_provider\n",
    "    # alias over the Device data provider\n",
    "    device_data_provider = hot3d_data_provider.device_data_provider\n",
    "    device_calibration = device_data_provider.get_device_calibration()\n",
    "    transform_device_cpf = device_calibration.get_transform_device_cpf()\n",
    "    # retrieve a list of timestamps for the sequence (in nanoseconds)\n",
    "    timestamps = device_data_provider.get_sequence_timestamps()\n",
    "    \n",
    "    # segment valid data\n",
    "    index = 0\n",
    "    valid_frames = 0\n",
    "    start_frames = []\n",
    "    end_frames = []\n",
    "    while(index<len(masks)):\n",
    "        value = masks[index]    \n",
    "        if value == True:\n",
    "            start = index\n",
    "            start_frames.append(start)\n",
    "            #print(\"start at: {}\".format(start))\n",
    "            while(value == True):\n",
    "                index += 1\n",
    "                if index<len(masks):\n",
    "                    value = masks[index]\n",
    "                else:\n",
    "                    break                \n",
    "            end = index-1\n",
    "            end_frames.append(end)\n",
    "            #print(\"end at: {}\".format(end))\n",
    "            valid_frames += end - start + 1            \n",
    "        else:\n",
    "            index += 1\n",
    "            \n",
    "    segment_num = len(start_frames)\n",
    "    local_time = time.asctime(time.localtime(time.time()))\n",
    "    print('\\nprocessing starts at ' + local_time)    \n",
    "    for i in range(segment_num):\n",
    "        start_frame = start_frames[i]\n",
    "        end_frame = end_frames[i]\n",
    "        frame_length = end_frame - start_frame + 1\n",
    "        if frame_length < valid_frame_length:\n",
    "            continue\n",
    "        print(\"start frame: {}, end frame: {}, length: {}\".format(start_frame, end_frame, frame_length))\n",
    "        \n",
    "        timestamps_data = np.zeros((frame_length, 1))\n",
    "        head_data = np.zeros((frame_length, 10)) # head_direction (3) + head_translation (3) + head_rotation (4, quat_xyzw)\n",
    "        gaze_data = np.zeros((frame_length, 6)) # gaze_direction (3) + gaze_center_in_world (3)        \n",
    "        hand_data = np.zeros((frame_length, 44)) # left_hand (22) + right_hand (22), hand = wrist_pose (7, translation (3) + rotation (4)) + joint_angles (15)            \n",
    "        hand_joint_data = np.zeros((frame_length, 122)) # left_hand (20*3) + right_hand (20*3) + attended_hand_gt + attended_hand_baseline (closest_hand) \n",
    "        hand_joint_initial_data = np.zeros((frame_length, 122)) # left_hand (20*3) + right_hand (20*3) + attended_hand_gt + closest_hand\n",
    "        object_data = np.zeros((frame_length, 48)) # 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.zeros((frame_length, 144)) # bounding box information: 6 objects (at most 6 objects) * 8 vertexes * 3\n",
    "        object_bbx_left_hand_data = np.zeros((frame_length, 144)) # bounding box information of the objects ranked using distances to the left hand\n",
    "        object_bbx_right_hand_data = np.zeros((frame_length, 144)) # bounding box information of the objects ranked using distances to the right hand\n",
    "        object_bbx_left_hand_initial_data = np.zeros((frame_length, 144)) # bounding box information of the objects ranked using distances to the left hand\n",
    "        object_bbx_right_hand_initial_data = np.zeros((frame_length, 144)) # bounding box information of the objects ranked using distances to the right hand\n",
    "        \n",
    "        # extract the valid frames\n",
    "        for frame in range(start_frame, end_frame+1):\n",
    "            timestamp_ns = timestamps[frame]\n",
    "            timestamps_data[frame-start_frame] = timestamp_ns\n",
    "            \n",
    "            # extract head data\n",
    "            headset_pose3d_with_dt = device_pose_provider.get_pose_at_timestamp(\n",
    "                                    timestamp_ns=timestamp_ns,\n",
    "                                    time_query_options=TimeQueryOptions.CLOSEST,\n",
    "                                    time_domain=TimeDomain.TIME_CODE)                    \n",
    "            headset_pose3d = headset_pose3d_with_dt.pose3d\n",
    "            T_world_device = headset_pose3d.T_world_device\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 = T_world_device @ transform_device_cpf            \n",
    "            head_translation = T_world_cpf.translation()[0]\n",
    "            head_center_in_cpf = np.array([0, 0, 1.0], dtype = np.float64)\n",
    "            head_center_in_world = T_world_cpf @ head_center_in_cpf\n",
    "            head_center_in_world = head_center_in_world.reshape(3, )\n",
    "            head_direction = head_center_in_world - head_translation\n",
    "            head_direction = np.array([x / np.linalg.norm(head_direction) for x in head_direction])            \n",
    "            head_rotation = np.roll(T_world_cpf.rotation().to_quat()[0], -1) # change from w,x,y,z to x,y,z,w\n",
    "            head_data[frame-start_frame, 0:3] = head_direction\n",
    "            head_data[frame-start_frame, 3:6] = head_translation\n",
    "            head_data[frame-start_frame, 6:10] = head_rotation\n",
    "            \n",
    "            # extract eye gaze data\n",
    "            aria_eye_gaze_data = device_data_provider.get_eye_gaze(timestamp_ns)                \n",
    "            yaw = aria_eye_gaze_data.yaw\n",
    "            pitch = aria_eye_gaze_data.pitch\n",
    "            depth = aria_eye_gaze_data.depth\n",
    "            if depth == 0:\n",
    "                depth = 1\n",
    "            gaze_center_in_cpf = np.array([tan(yaw), tan(pitch), 1.0], dtype = np.float64)*depth\n",
    "            gaze_center_in_world = T_world_cpf @ gaze_center_in_cpf\n",
    "            gaze_center_in_world = gaze_center_in_world.reshape(3, )\n",
    "            gaze_direction = gaze_center_in_world - head_translation\n",
    "            gaze_direction = np.array([x / np.linalg.norm(gaze_direction) for x in gaze_direction])\n",
    "            # in rare cases, yaw, pitch is nan\n",
    "            if np.isnan(np.sum(gaze_direction)):\n",
    "                # use previous frame as an alternative\n",
    "                gaze_direction = gaze_data[frame-start_frame-1, 0:3]\n",
    "                gaze_center_in_world = gaze_data[frame-start_frame-1, 3:6]\n",
    "            gaze_data[frame-start_frame, 0:3] = gaze_direction\n",
    "            gaze_data[frame-start_frame, 3:6] = gaze_center_in_world\n",
    "            \n",
    "            # extract hand data\n",
    "            hand_poses_with_dt = hand_data_provider.get_pose_at_timestamp(\n",
    "                                timestamp_ns=timestamp_ns,\n",
    "                                time_query_options=TimeQueryOptions.CLOSEST,\n",
    "                                time_domain=TimeDomain.TIME_CODE)            \n",
    "            hand_pose_collection = hand_poses_with_dt.pose3d_collection\n",
    "            left_hand = hand_pose_collection.poses[Handedness.Left]\n",
    "            left_hand_translation = left_hand.wrist_pose.translation()[0]\n",
    "            left_hand_rotation = np.roll(left_hand.wrist_pose.rotation().to_quat()[0], -1) # change from w,x,y,z to x,y,z,w\n",
    "            left_hand_joint_angles = left_hand.joint_angles\n",
    "            left_hand_joints = hand_data_provider.get_hand_landmarks(left_hand).numpy().reshape(-1)\n",
    "            left_hand_initial = HandPose(Handedness.Left, left_hand.wrist_pose, np.zeros(15))\n",
    "            left_hand_initial_joints = hand_data_provider.get_hand_landmarks(left_hand_initial).numpy().reshape(-1)            \n",
    "            right_hand = hand_pose_collection.poses[Handedness.Right]\n",
    "            right_hand_translation = right_hand.wrist_pose.translation()[0]\n",
    "            right_hand_rotation = np.roll(right_hand.wrist_pose.rotation().to_quat()[0], -1) # change from w,x,y,z to x,y,z,w\n",
    "            right_hand_joint_angles = right_hand.joint_angles\n",
    "            right_hand_joints = hand_data_provider.get_hand_landmarks(right_hand).numpy().reshape(-1)\n",
    "            right_hand_initial = HandPose(Handedness.Right, right_hand.wrist_pose, np.zeros(15))\n",
    "            right_hand_initial_joints = hand_data_provider.get_hand_landmarks(right_hand_initial).numpy().reshape(-1)\n",
    "            \n",
    "            left_hand_direction = np.mean(left_hand_joints.reshape((20, 3)), axis=0) - head_translation\n",
    "            left_hand_direction = np.array([x / np.linalg.norm(left_hand_direction) for x in left_hand_direction]) \n",
    "            left_hand_distance_to_gaze = np.arccos(np.sum(gaze_direction*left_hand_direction))\n",
    "            right_hand_direction = np.mean(right_hand_joints.reshape((20, 3)), axis=0) - head_translation\n",
    "            right_hand_direction = np.array([x / np.linalg.norm(right_hand_direction) for x in right_hand_direction]) \n",
    "            right_hand_distance_to_gaze = np.arccos(np.sum(gaze_direction*right_hand_direction))\n",
    "            if left_hand_distance_to_gaze < right_hand_distance_to_gaze:\n",
    "                hand_joint_data[frame-start_frame, 120:121] = 0\n",
    "            else:\n",
    "                hand_joint_data[frame-start_frame, 120:121] = 1\n",
    "\n",
    "            left_hand_initial_direction = np.mean(left_hand_initial_joints.reshape((20, 3)), axis=0) - head_translation\n",
    "            left_hand_initial_direction = np.array([x / np.linalg.norm(left_hand_initial_direction) for x in left_hand_initial_direction]) \n",
    "            left_hand_initial_distance_to_gaze = np.arccos(np.sum(gaze_direction*left_hand_initial_direction))\n",
    "            right_hand_initial_direction = np.mean(right_hand_initial_joints.reshape((20, 3)), axis=0) - head_translation\n",
    "            right_hand_initial_direction = np.array([x / np.linalg.norm(right_hand_initial_direction) for x in right_hand_initial_direction]) \n",
    "            right_hand_initial_distance_to_gaze = np.arccos(np.sum(gaze_direction*right_hand_initial_direction))\n",
    "            if left_hand_initial_distance_to_gaze < right_hand_initial_distance_to_gaze:\n",
    "                hand_joint_initial_data[frame-start_frame, 120:121] = 0\n",
    "            else:\n",
    "                hand_joint_initial_data[frame-start_frame, 120:121] = 1\n",
    "            \n",
    "            hand_data[frame-start_frame, 0:3] = left_hand_translation\n",
    "            hand_data[frame-start_frame, 3:7] = left_hand_rotation\n",
    "            hand_data[frame-start_frame, 7:22] = left_hand_joint_angles\n",
    "            hand_data[frame-start_frame, 22:25] = right_hand_translation\n",
    "            hand_data[frame-start_frame, 25:29] = right_hand_rotation\n",
    "            hand_data[frame-start_frame, 29:44] = right_hand_joint_angles\n",
    "            hand_joint_data[frame-start_frame, 0:60] = left_hand_joints\n",
    "            hand_joint_data[frame-start_frame, 60:120] = right_hand_joints\n",
    "            hand_joint_initial_data[frame-start_frame, 0:60] = left_hand_initial_joints\n",
    "            hand_joint_initial_data[frame-start_frame, 60:120] = right_hand_initial_joints\n",
    "                        \n",
    "            # extract object data\n",
    "            object_poses_with_dt = object_pose_data_provider.get_pose_at_timestamp(\n",
    "                                timestamp_ns=timestamp_ns,\n",
    "                                time_query_options=TimeQueryOptions.CLOSEST,\n",
    "                                time_domain=TimeDomain.TIME_CODE)\n",
    "            objects_pose3d = object_poses_with_dt.pose3d_collection.poses\n",
    "            object_num = len(objects_pose3d)\n",
    "            objects_distance_to_left_hand = {}\n",
    "            objects_distance_to_right_hand = {}            \n",
    "            objects_distance_to_left_hand_initial = {}\n",
    "            objects_distance_to_right_hand_initial = {}\n",
    "            objects_pose3d_dict = {}\n",
    "            item = 0\n",
    "            for (object_uid, object_pose3d) in objects_pose3d.items():\n",
    "                object_translation = object_pose3d.T_world_object.translation()[0]                \n",
    "                object_distance_to_left_hand = np.mean(np.linalg.norm(left_hand_joints.reshape((20, 3))-object_translation, axis=1))\n",
    "                object_distance_to_right_hand = np.mean(np.linalg.norm(right_hand_joints.reshape((20, 3))-object_translation, axis=1))\n",
    "                object_distance_to_left_hand_initial = np.mean(np.linalg.norm(left_hand_initial_joints.reshape((20, 3))-object_translation, axis=1))\n",
    "                object_distance_to_right_hand_initial = np.mean(np.linalg.norm(right_hand_initial_joints.reshape((20, 3))-object_translation, axis=1))                \n",
    "                objects_distance_to_left_hand[object_uid] = object_distance_to_left_hand              \n",
    "                objects_distance_to_right_hand[object_uid] = object_distance_to_right_hand\n",
    "                objects_distance_to_left_hand_initial[object_uid] = object_distance_to_left_hand_initial\n",
    "                objects_distance_to_right_hand_initial[object_uid] = object_distance_to_right_hand_initial\n",
    "                objects_pose3d_dict[object_uid] = object_pose3d.T_world_object                                \n",
    "                item += 1\n",
    "                \n",
    "            objects_distance_to_left_hand_sorted = sorted(objects_distance_to_left_hand.items(), key = lambda kv:(kv[1], kv[0]))\n",
    "            objects_distance_to_right_hand_sorted = sorted(objects_distance_to_right_hand.items(), key = lambda kv:(kv[1], kv[0]))            \n",
    "            left_object_closest_uid = objects_distance_to_left_hand_sorted[0][0]\n",
    "            left_object_closest_distance = objects_distance_to_left_hand_sorted[0][1]\n",
    "            right_object_closest_uid = objects_distance_to_right_hand_sorted[0][0]\n",
    "            right_object_closest_distance = objects_distance_to_right_hand_sorted[0][1]\n",
    "            if left_object_closest_distance < right_object_closest_distance:\n",
    "                hand_joint_data[frame-start_frame, -1] = 0\n",
    "            else:\n",
    "                hand_joint_data[frame-start_frame, -1] = 1\n",
    "\n",
    "            objects_distance_to_left_hand_initial_sorted = sorted(objects_distance_to_left_hand_initial.items(), key = lambda kv:(kv[1], kv[0]))\n",
    "            objects_distance_to_right_hand_initial_sorted = sorted(objects_distance_to_right_hand_initial.items(), key = lambda kv:(kv[1], kv[0]))\n",
    "            left_initial_object_closest_uid = objects_distance_to_left_hand_initial_sorted[0][0]\n",
    "            left_initial_object_closest_distance = objects_distance_to_left_hand_initial_sorted[0][1]\n",
    "            right_initial_object_closest_uid = objects_distance_to_right_hand_initial_sorted[0][0]\n",
    "            right_initial_object_closest_distance = objects_distance_to_right_hand_initial_sorted[0][1]\n",
    "            if left_initial_object_closest_distance < right_initial_object_closest_distance:\n",
    "                hand_joint_initial_data[frame-start_frame, -1] = 0\n",
    "            else:\n",
    "                hand_joint_initial_data[frame-start_frame, -1] = 1\n",
    "            \n",
    "            item = 0\n",
    "            for object_uid in objects_pose3d_dict:\n",
    "                object_pose3d = objects_pose3d_dict[object_uid]\n",
    "                object_translation = object_pose3d.translation()[0]\n",
    "                object_rotation = np.roll(object_pose3d.rotation().to_quat()[0], -1) # change from w,x,y,z to x,y,z,w                \n",
    "                object_data[frame-start_frame, item*8:item*8+1] = object_uid\n",
    "                object_data[frame-start_frame, item*8+1:item*8+4] = object_translation\n",
    "                object_data[frame-start_frame, item*8+4:item*8+8] = object_rotation\n",
    "                bbx = object_bbx[object_uid]\n",
    "                #print(\"uid: {}, bbx: {}\".format(object_uid, bbx))\n",
    "                x_min = bbx[0]\n",
    "                x_max = bbx[1]\n",
    "                y_min = bbx[2]\n",
    "                y_max = bbx[3]\n",
    "                z_min = bbx[4]\n",
    "                z_max = bbx[5]\n",
    "                bbx_vertex = np.array([x_min, y_min, z_min], dtype = np.float64)                                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )                \n",
    "                object_bbx_data[frame-start_frame, item*24:item*24+3] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_min, z_min], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_data[frame-start_frame, item*24+3:item*24+6] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_min, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_data[frame-start_frame, item*24+6:item*24+9] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_min, y_min, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_data[frame-start_frame, item*24+9:item*24+12] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_min, y_max, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_data[frame-start_frame, item*24+12:item*24+15] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_max, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_data[frame-start_frame, item*24+15:item*24+18] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_max, z_min], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_data[frame-start_frame, item*24+18:item*24+21] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_min, y_max, z_min], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_data[frame-start_frame, item*24+21:item*24+24] = bbx_vertex                \n",
    "                item += 1\n",
    "                \n",
    "            for item in range(len(objects_distance_to_left_hand_sorted)):\n",
    "                object_uid = objects_distance_to_left_hand_sorted[item][0]\n",
    "                object_pose3d = objects_pose3d_dict[object_uid]            \n",
    "                object_translation = object_pose3d.translation()[0]\n",
    "                object_rotation = np.roll(object_pose3d.rotation().to_quat()[0], -1) # change from w,x,y,z to x,y,z,w                \n",
    "                bbx = object_bbx[object_uid]\n",
    "                #print(\"uid: {}, bbx: {}\".format(object_uid, bbx))\n",
    "                x_min = bbx[0]\n",
    "                x_max = bbx[1]\n",
    "                y_min = bbx[2]\n",
    "                y_max = bbx[3]\n",
    "                z_min = bbx[4]\n",
    "                z_max = bbx[5]\n",
    "                bbx_vertex = np.array([x_min, y_min, z_min], dtype = np.float64)                                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )                \n",
    "                object_bbx_left_hand_data[frame-start_frame, item*24:item*24+3] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_min, z_min], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_left_hand_data[frame-start_frame, item*24+3:item*24+6] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_min, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_left_hand_data[frame-start_frame, item*24+6:item*24+9] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_min, y_min, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_left_hand_data[frame-start_frame, item*24+9:item*24+12] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_min, y_max, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_left_hand_data[frame-start_frame, item*24+12:item*24+15] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_max, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_left_hand_data[frame-start_frame, item*24+15:item*24+18] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_max, z_min], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_left_hand_data[frame-start_frame, item*24+18:item*24+21] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_min, y_max, z_min], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_left_hand_data[frame-start_frame, item*24+21:item*24+24] = bbx_vertex                \n",
    "\n",
    "            for item in range(len(objects_distance_to_right_hand_sorted)):\n",
    "                object_uid = objects_distance_to_right_hand_sorted[item][0]\n",
    "                object_pose3d = objects_pose3d_dict[object_uid]            \n",
    "                object_translation = object_pose3d.translation()[0]\n",
    "                object_rotation = np.roll(object_pose3d.rotation().to_quat()[0], -1) # change from w,x,y,z to x,y,z,w                \n",
    "                bbx = object_bbx[object_uid]\n",
    "                #print(\"uid: {}, bbx: {}\".format(object_uid, bbx))\n",
    "                x_min = bbx[0]\n",
    "                x_max = bbx[1]\n",
    "                y_min = bbx[2]\n",
    "                y_max = bbx[3]\n",
    "                z_min = bbx[4]\n",
    "                z_max = bbx[5]\n",
    "                bbx_vertex = np.array([x_min, y_min, z_min], dtype = np.float64)                                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )                \n",
    "                object_bbx_right_hand_data[frame-start_frame, item*24:item*24+3] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_min, z_min], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_right_hand_data[frame-start_frame, item*24+3:item*24+6] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_min, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_right_hand_data[frame-start_frame, item*24+6:item*24+9] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_min, y_min, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_right_hand_data[frame-start_frame, item*24+9:item*24+12] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_min, y_max, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_right_hand_data[frame-start_frame, item*24+12:item*24+15] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_max, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_right_hand_data[frame-start_frame, item*24+15:item*24+18] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_max, z_min], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_right_hand_data[frame-start_frame, item*24+18:item*24+21] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_min, y_max, z_min], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_right_hand_data[frame-start_frame, item*24+21:item*24+24] = bbx_vertex                \n",
    "\n",
    "            for item in range(len(objects_distance_to_left_hand_initial_sorted)):\n",
    "                object_uid = objects_distance_to_left_hand_initial_sorted[item][0]\n",
    "                object_pose3d = objects_pose3d_dict[object_uid]            \n",
    "                object_translation = object_pose3d.translation()[0]\n",
    "                object_rotation = np.roll(object_pose3d.rotation().to_quat()[0], -1) # change from w,x,y,z to x,y,z,w                \n",
    "                bbx = object_bbx[object_uid]\n",
    "                #print(\"uid: {}, bbx: {}\".format(object_uid, bbx))\n",
    "                x_min = bbx[0]\n",
    "                x_max = bbx[1]\n",
    "                y_min = bbx[2]\n",
    "                y_max = bbx[3]\n",
    "                z_min = bbx[4]\n",
    "                z_max = bbx[5]\n",
    "                bbx_vertex = np.array([x_min, y_min, z_min], dtype = np.float64)                                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )                \n",
    "                object_bbx_left_hand_initial_data[frame-start_frame, item*24:item*24+3] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_min, z_min], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_left_hand_initial_data[frame-start_frame, item*24+3:item*24+6] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_min, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_left_hand_initial_data[frame-start_frame, item*24+6:item*24+9] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_min, y_min, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_left_hand_initial_data[frame-start_frame, item*24+9:item*24+12] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_min, y_max, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_left_hand_initial_data[frame-start_frame, item*24+12:item*24+15] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_max, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_left_hand_initial_data[frame-start_frame, item*24+15:item*24+18] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_max, z_min], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_left_hand_initial_data[frame-start_frame, item*24+18:item*24+21] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_min, y_max, z_min], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_left_hand_initial_data[frame-start_frame, item*24+21:item*24+24] = bbx_vertex                \n",
    "\n",
    "            for item in range(len(objects_distance_to_right_hand_initial_sorted)):\n",
    "                object_uid = objects_distance_to_right_hand_initial_sorted[item][0]\n",
    "                object_pose3d = objects_pose3d_dict[object_uid]            \n",
    "                object_translation = object_pose3d.translation()[0]\n",
    "                object_rotation = np.roll(object_pose3d.rotation().to_quat()[0], -1) # change from w,x,y,z to x,y,z,w                \n",
    "                bbx = object_bbx[object_uid]\n",
    "                #print(\"uid: {}, bbx: {}\".format(object_uid, bbx))\n",
    "                x_min = bbx[0]\n",
    "                x_max = bbx[1]\n",
    "                y_min = bbx[2]\n",
    "                y_max = bbx[3]\n",
    "                z_min = bbx[4]\n",
    "                z_max = bbx[5]\n",
    "                bbx_vertex = np.array([x_min, y_min, z_min], dtype = np.float64)                                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )                \n",
    "                object_bbx_right_hand_initial_data[frame-start_frame, item*24:item*24+3] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_min, z_min], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_right_hand_initial_data[frame-start_frame, item*24+3:item*24+6] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_min, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_right_hand_initial_data[frame-start_frame, item*24+6:item*24+9] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_min, y_min, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_right_hand_initial_data[frame-start_frame, item*24+9:item*24+12] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_min, y_max, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_right_hand_initial_data[frame-start_frame, item*24+12:item*24+15] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_max, z_max], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_right_hand_initial_data[frame-start_frame, item*24+15:item*24+18] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_max, y_max, z_min], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_right_hand_initial_data[frame-start_frame, item*24+18:item*24+21] = bbx_vertex\n",
    "                bbx_vertex = np.array([x_min, y_max, z_min], dtype = np.float64)                \n",
    "                bbx_vertex = (object_pose3d @ bbx_vertex).reshape(3, )\n",
    "                object_bbx_right_hand_initial_data[frame-start_frame, item*24+21:item*24+24] = bbx_vertex                \n",
    "        \n",
    "        # save the data\n",
    "        timestamps_path = save_path + str(start_frame) + \"_\" + str(end_frame) + '_timestamps.npy'\n",
    "        head_path = save_path + str(start_frame) + \"_\" + str(end_frame) + '_head.npy'\n",
    "        gaze_path = save_path + str(start_frame) + \"_\" + str(end_frame) + '_gaze.npy'\n",
    "        hand_path = save_path + str(start_frame) + \"_\" + str(end_frame) + '_hand.npy'        \n",
    "        hand_joint_path = save_path + str(start_frame) + \"_\" + str(end_frame) + '_handjoints.npy'\n",
    "        hand_joint_initial_path = save_path + str(start_frame) + \"_\" + str(end_frame) + '_inithandjoints.npy'\n",
    "        object_path = save_path + str(start_frame) + \"_\" + str(end_frame) + '_objects.npy'\n",
    "        object_bbx_path = save_path + str(start_frame) + \"_\" + str(end_frame) + '_object_bbx.npy'\n",
    "        object_bbx_left_hand_path = save_path + str(start_frame) + \"_\" + str(end_frame) + '_object_bbxleft.npy'\n",
    "        object_bbx_right_hand_path = save_path + str(start_frame) + \"_\" + str(end_frame) + '_object_bbxright.npy'\n",
    "        object_bbx_left_hand_initial_path = save_path + str(start_frame) + \"_\" + str(end_frame) + '_object_initbbxleft.npy'\n",
    "        object_bbx_right_hand_initial_path = save_path + str(start_frame) + \"_\" + str(end_frame) + '_object_initbbxright.npy'\n",
    "        \n",
    "        np.save(timestamps_path, timestamps_data)\n",
    "        np.save(head_path, head_data)\n",
    "        np.save(gaze_path, gaze_data)\n",
    "        np.save(hand_path, hand_data)\n",
    "        np.save(hand_joint_path, hand_joint_data)\n",
    "        np.save(hand_joint_initial_path, hand_joint_initial_data)\n",
    "        np.save(object_path, object_data)\n",
    "        np.save(object_bbx_path, object_bbx_data)\n",
    "        np.save(object_bbx_left_hand_path, object_bbx_left_hand_data)\n",
    "        np.save(object_bbx_right_hand_path, object_bbx_right_hand_data)\n",
    "        np.save(object_bbx_left_hand_initial_path, object_bbx_left_hand_initial_data)\n",
    "        np.save(object_bbx_right_hand_initial_path, object_bbx_right_hand_initial_data)\n",
    "        \n",
    "    local_time = time.asctime(time.localtime(time.time()))\n",
    "    print('\\nprocessing ends at ' + local_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec30e81d-8812-4665-be58-00a7e0aa1915",
   "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
}
