{ "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