import os import sys sys.path.append('.') import pickle import numpy as np # import open3d as o3d import torch # import json from tqdm import tqdm import traceback from manopth.manolayer import ManoLayer # from utils.hoi_io2 import get_num_frame_v2 from utils.organize_dataset import load_sequence_names_from_organized_record, txt2intrinsic, load_simplied_nokov_objs_mesh, load_organized_mano_info, load_dates_from_organized_record def load_objs_orientation(root, video_id, frame_list = None): ''' return: tool_pose_batch, obj_pose_batch ''' date = video_id[:8] original_nokov_data_dir = os.path.join(root, date, video_id, 'nokov') assert os.path.exists(original_nokov_data_dir) nokov_data_filenames = os.listdir(original_nokov_data_dir) if frame_list is not None: frame_idx_list = [(int(frame)-1) for frame in frame_list] tool_id = nokov_data_filenames[0].split("_")[1] # obj1 obj_id = nokov_data_filenames[0].split("_")[2] # obj2 tool_pose_path = os.path.join(root, 'HO_poses', date, video_id, 'objpose', f'{tool_id}.npy') assert os.path.exists(tool_pose_path) if frame_list is None: tool_pose_batch = torch.from_numpy(np.load(tool_pose_path)) else: tool_pose_batch = torch.from_numpy(np.load(tool_pose_path))[frame_idx_list, ...] obj_pose_path = os.path.join(root, 'HO_poses', date, video_id, 'objpose', f'{obj_id}.npy') assert os.path.exists(obj_pose_path) if frame_list is None: obj_pose_batch = torch.from_numpy(np.load(obj_pose_path)) else: obj_pose_batch = torch.from_numpy(np.load(obj_pose_path))[frame_idx_list, ...] return tool_pose_batch, obj_pose_batch if __name__ == '__main__': # date_list = ['20230930'] device = 'cuda:0' root = '/data3/hlyang/results' upload_root = '/data3/hlyang/HOI-mocap' nokov_root = upload_root dataset_root = os.path.join(root, 'dataset') hand_pose_organized_record_path = os.path.join(dataset_root, 'organized_record.txt') date_list = load_dates_from_organized_record(hand_pose_organized_record_path) # save_root = '/data3/hlyang/results/test_data' save_root = '/data3/datasets/xueyi/taco/processed_data' mano_path = "/data1/xueyi/mano_models/mano/models" for date in date_list: video_list = load_sequence_names_from_organized_record(hand_pose_organized_record_path, date) for video_id in tqdm(video_list): try: save_dir = os.path.join(save_root, date) os.makedirs(save_dir, exist_ok=True) (right_hand_pose, right_hand_trans, right_hand_shape) = load_organized_mano_info(dataset_root, date, video_id, frame_list = None, right_hand_bool=True) (left_hand_pose, left_hand_trans, left_hand_shape) = load_organized_mano_info(dataset_root, date, video_id, frame_list = None, right_hand_bool=False) tool_verts_batch, tool_faces, obj_verts_batch, obj_faces = load_simplied_nokov_objs_mesh(nokov_root, video_id, frame_list = None, use_cm = True) tool_pose_batch, obj_pose_batch = load_objs_orientation(nokov_root, video_id, frame_list = None) batch_size = right_hand_pose.shape[0] if right_hand_shape is not None and left_hand_shape is not None: use_shape = True else: use_shape = False right_hand_pose = right_hand_pose.to(device) right_hand_trans = right_hand_trans.to(device) left_hand_pose = left_hand_pose.to(device) left_hand_trans = left_hand_trans.to(device) tool_verts_batch = tool_verts_batch.to(device) obj_verts_batch = obj_verts_batch.to(device) if use_shape: right_hand_shape = right_hand_shape.to(device) left_hand_shape = left_hand_shape.to(device) # get hand meshes use_pca = False ncomps = 45 left_hand_mano_layer = ManoLayer(mano_root=mano_path, use_pca=use_pca, ncomps=ncomps, side='left', center_idx = 0).to(device) right_hand_mano_layer = ManoLayer(mano_root=mano_path, use_pca=use_pca, ncomps=ncomps, side='right', center_idx = 0).to(device) with torch.no_grad(): if use_shape: right_hand_verts, _ = right_hand_mano_layer(right_hand_pose, right_hand_shape.unsqueeze(0).expand(batch_size, -1)) else: right_hand_verts, _ = right_hand_mano_layer(right_hand_pose) right_hand_verts = right_hand_verts right_hand_verts = right_hand_verts / 1000.0 right_hand_verts += right_hand_trans.unsqueeze(1) if use_shape: left_hand_verts, _ = left_hand_mano_layer(left_hand_pose, left_hand_shape.unsqueeze(0).expand(batch_size, -1)) else: left_hand_verts, _ = left_hand_mano_layer(left_hand_pose) left_hand_verts = left_hand_verts left_hand_verts = left_hand_verts / 1000.0 left_hand_verts += left_hand_trans.unsqueeze(1) # left hand hand_pose = left_hand_pose.cpu().numpy() hand_trans = left_hand_trans.cpu().numpy() if use_shape: hand_shape = left_hand_shape.cpu().numpy() else: hand_shape = torch.zeros(10).cpu().numpy() hand_verts = left_hand_verts.cpu().numpy() hand_faces = left_hand_mano_layer.th_faces.cpu().numpy() obj_verts = obj_verts_batch.cpu().numpy() obj_faces = obj_faces.cpu().numpy() # rgt_hand_pose = right_hand_pose.cpu().numpy() rgt_hand_trans = right_hand_trans.cpu().numpy() if use_shape: rgt_hand_shape = right_hand_shape.cpu().numpy() else: rgt_hand_shape = torch.zeros(10).cpu().numpy() rgt_hand_verts = right_hand_verts.cpu().numpy() rgt_hand_faces = right_hand_mano_layer.th_faces.cpu().numpy() rgt_obj_verts = tool_verts_batch.cpu().numpy() rgt_obj_faces = tool_faces.cpu().numpy() # print(hand_verts.shape, hand_faces.shape, obj_verts.shape, obj_faces.shape) save_path = os.path.join(save_dir, f'right_{video_id}.pkl') # save_data = { # 'hand_pose': hand_pose, 'hand_trans': hand_trans, 'hand_shape': hand_shape, 'hand_verts': hand_verts, 'hand_faces': hand_faces, 'obj_verts': obj_verts, 'obj_faces': obj_faces, 'obj_pose': obj_pose_batch.cpu().numpy(), # } save_data = { 'hand_pose': rgt_hand_pose, 'hand_trans': rgt_hand_trans, 'hand_shape': rgt_hand_shape, 'hand_verts': rgt_hand_verts, 'hand_faces': rgt_hand_faces, 'obj_verts': rgt_obj_verts, 'obj_faces': rgt_obj_faces, 'obj_pose': tool_pose_batch.cpu().numpy(), } # print(save_data['hand_pose'].shape, save_data['hand_trans'].shape, save_data['hand_shape'].shape, save_data['hand_verts'].shape, save_data['hand_faces'].shape, save_data['obj_verts'].shape, save_data['obj_faces'].shape, save_data['obj_pose'].shape) pickle.dump(save_data, open(save_path, 'wb')) except Exception as err: traceback.print_exc() print(err) continue