Spaces:
Sleeping
Sleeping
File size: 7,957 Bytes
710e818 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 |
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 |