Spaces:
Runtime error
Runtime error
import os | |
import pdb | |
import pickle as pkl | |
from torch.utils.data import Dataset, DataLoader | |
# from omegaconf import DictConfig, OmegaConf | |
from scipy.spatial.transform import Rotation as R | |
from models.joints_to_smplx import joints_to_smpl, JointsToSMPLX | |
from utils import * | |
from constants import * | |
from datasets.trumans import TrumansDataset | |
from models.synhsi import Unet | |
# from hydra import compose, initialize | |
import yaml | |
ACT_TYPE = ['scene', 'grasp', 'artic', 'none'] | |
def convert_trajectory(trajectory): | |
trajectory_new = [[t['x'], t['y'], t['z']] for t in trajectory] | |
trajectory_new = np.array(trajectory_new) | |
return trajectory_new | |
def get_base_speed(cfg, trajectory, is_zup=True): | |
trajectory_layer = trajectory | |
if is_zup: | |
trajectory_layer = zup_to_yup(trajectory_layer) | |
trajectory2D = trajectory_layer[:, [0, 2]] | |
distance = np.sum(np.linalg.norm(trajectory2D[1:] - trajectory2D[:-1], axis=1)) | |
speed = trajectory_layer.shape[0] // distance | |
print('Base Speed:', speed, flush=True) | |
return speed | |
def get_guidance(cfg, trajectory, samplers, act_type='none', speed=35): | |
trajectory_layer = trajectory | |
# trajectory_layer = zup_to_yup(trajectory_layer) | |
print(trajectory_layer.shape) | |
if cfg.action_type != 'pure_inter': | |
#TODO | |
midpoints = trajectory_layer[[0] * cfg.len_pre + list(range(0, len(trajectory_layer), speed)) + [-1] * (cfg.len_act + (1 if cfg.stay_and_act else 0))] | |
# midpoints[0, 0] = 1.4164 | |
# midpoints[0, 2] = 2.2544 | |
# midpoints = trajectory_layer[[0] * cfg.len_pre + [25, 50] + [50] + [70, 90]] | |
else: | |
midpoints = trajectory_layer[[0] * cfg.len_pre + [0] + [0] * (cfg.len_act + 1 if cfg.stay_and_act else 0)] | |
midpoints = torch.tensor(midpoints).float().to(cfg.device) | |
max_step = midpoints.shape[0] - 1 | |
mat_init = cfg.batch_size * [np.eye(4)] | |
mat_init = torch.from_numpy(np.stack(mat_init, axis=0)).float().to(cfg.device) | |
print(midpoints) | |
mat_init[:, 0, 3] = midpoints[0, 0] | |
mat_init[:, 2, 3] = midpoints[0, 2] | |
dx = midpoints[cfg.len_pre + 1, 0] - midpoints[0, 0] | |
dz = midpoints[cfg.len_pre + 1, 2] - midpoints[0, 2] | |
print(-np.arctan2(dx.item(), dz.item()), dx, dz) | |
mat_rot_y = R.from_rotvec(np.array([0, np.arctan2(dx.item(), dz.item()), 0])).as_matrix() | |
# mat_rot_y = R.from_rotvec(np.array([0, np.arctan2(-1, -2), 0])).as_matrix() | |
mat_init[:, :3, :3] = torch.from_numpy(mat_rot_y).float().to(cfg.device) | |
# goal_list = torch.zeros((max_step, cfg.batch_size, cfg.dataset.seq_len, 3)).float().to(cfg.device) | |
goal_list = [] | |
action_label_list = [] | |
# action_label_list = torch.zeros((max_step, cfg.batch_size, cfg.dataset.seq_len, cfg.dataset.nb_actions)).float().to(cfg.device) | |
sampler_list = [] | |
if act_type == 'none': | |
for s in range(max_step): | |
goal = torch.zeros((cfg.batch_size, 1, 3)).float().to(cfg.device) | |
goal[:, :] = midpoints[s + 1] | |
goal_list.append(goal) | |
if cfg.dataset.nb_actions > 0: | |
action_label = torch.zeros((cfg.batch_size, cfg.dataset.seq_len, cfg.dataset.nb_actions)).float().to(cfg.device) | |
action_label_list.append(action_label) | |
else: | |
action_label_list.append(None) | |
sampler_list.append(samplers['body']) | |
elif act_type == 'write': | |
midpoints = torch.from_numpy(trajectory_layer[::4]).float().to(cfg.device) | |
for s in range(midpoints.shape[0] // 16): | |
goal = torch.zeros((cfg.batch_size, 16, 3)).float().to(cfg.device) | |
goal[:, :] = midpoints[s * 16: (s + 1) * 16] | |
goal_list.append(goal) | |
if cfg.dataset.nb_actions > 0: | |
action_label = torch.zeros((cfg.batch_size, cfg.dataset.seq_len, cfg.dataset.nb_actions)).float().to(cfg.device) | |
action_label_list.append(action_label) | |
else: | |
action_label_list.append(None) | |
sampler_list.append(samplers['hand']) | |
elif act_type == 'scene': | |
for s in range(max_step): | |
goal = torch.zeros((cfg.batch_size, 1, 3)).float().to(cfg.device) | |
goal[:, :] = midpoints[s + 1] | |
goal_list.append(goal) | |
sampler_list.append(samplers['body']) | |
action_label = torch.zeros((cfg.batch_size, cfg.dataset.seq_len, cfg.dataset.nb_actions)).float().to(cfg.device) | |
if s > max_step - cfg.len_act: | |
action_label[:, :, cfg.action_id] = 1. | |
action_label_list.append(action_label) | |
elif act_type == 'grasp': | |
for s in range(max_step): | |
if s != max_step - cfg.len_act: | |
goal = torch.zeros((cfg.batch_size, 1, 3)).float().to(cfg.device) | |
goal[:, :] = midpoints[s + 1] | |
goal_list.append(goal) | |
else: | |
# grasp_goal = zup_to_yup(np.array([[-0.32, -3.36, 0.395], | |
# [-0.3, -3.2, 0.395], | |
# [-0.25, -3, 0.394], | |
# [-0.147, -3.0, 0.395]])).reshape((cfg.batch_size, -1, 3)) | |
grasp_goal = zup_to_yup(np.array(trajectory['Object'])).reshape((cfg.batch_size, 1, 3)) | |
goal = torch.zeros((cfg.batch_size, 3, 3)).float().to(cfg.device) | |
goal[:, :] = torch.from_numpy(grasp_goal).float().to(cfg.device) | |
goal_list.append(goal) | |
action_label = torch.zeros((cfg.batch_size, cfg.dataset.seq_len, cfg.dataset.nb_actions)).float().to(cfg.device) | |
if s < max_step - cfg.len_act: | |
sampler_list.append(samplers['body']) | |
elif s == max_step - cfg.len_act: | |
sampler_list.append(samplers['hand']) | |
else: | |
sampler_list.append(samplers['body']) | |
if cfg.action_id != -1: | |
action_label[:, :, cfg.action_id] = 1. | |
action_label_list.append(action_label) | |
elif act_type == 'pure_inter': | |
sampler_list += [samplers['body']] * (cfg.len_pre + cfg.len_act) | |
goal = torch.zeros((cfg.batch_size, 1, 3)).float().to(cfg.device) | |
goal[:, :] = midpoints[0] | |
goal_list += [goal] * (cfg.len_pre + + cfg.len_act) | |
action_label = torch.zeros((cfg.batch_size, cfg.dataset.seq_len, cfg.dataset.nb_actions)).float().to(cfg.device) | |
action_label_list += [action_label.clone()] * cfg.len_pre | |
action_label[:, :, cfg.action_id] = 2. | |
action_label_list += [action_label.clone()] * cfg.len_act | |
return mat_init, goal_list, action_label_list, sampler_list | |
def sample_step(cfg, mat, obj_locs, goal_list, action_label_list, sampler_list): | |
max_step = len(goal_list) | |
fixed_points = None | |
fixed_frame = 2 | |
points_all = [] | |
cnt_fixed_frame = 0 | |
cnt_seq_len = 0 | |
for s in range(max_step): | |
print('step', s) | |
sampler = sampler_list[s] | |
if s != 0: | |
fixed_points = sampler.dataset.normalize_torch(transform_points(fixed_points, torch.inverse(mat))) | |
else: | |
if cfg.continue_last: | |
method_id = cfg.method_name.split('_')[-1] | |
method_name_last = cfg.method_name[:-1] + str(int(method_id) - 1) | |
mat = torch.from_numpy(np.load(os.path.join(cfg.exp_dir, f'{method_name_last}_mat.npy'))).to(sampler.device) | |
fixed_points = torch.from_numpy(np.load(os.path.join(cfg.exp_dir, f'{method_name_last}_fixed_points.npy'))).to(sampler.device) | |
fixed_points = sampler.dataset.normalize_torch(transform_points(fixed_points, torch.inverse(mat))) | |
samples, occs = sampler.p_sample_loop(fixed_points, obj_locs, mat, cfg.scene_name, goal_list[s], action_label_list[s]) | |
if 0 <= s < cfg.len_pre: | |
cnt_fixed_frame += sampler.fixed_frame | |
if 0 <= s < cfg.len_pre: | |
cnt_seq_len += cfg.dataset.seq_len | |
points_gene = samples[-1] | |
points_gene_np = points_gene.reshape(cfg.batch_size, cfg.dataset.seq_len, -1, 3).cpu().numpy() | |
if s == 0 or fixed_frame == 0: | |
#TODO | |
points_all.append(points_gene_np[:, fixed_frame - 1:]) | |
elif fixed_frame > 0: | |
points_all.append(points_gene_np[:, fixed_frame:]) | |
# fixed_frame = 0 if s == max_step - 1 else sampler_list[s + 1].fixed_frame | |
fixed_frame = sampler_list[s].fixed_frame if s == max_step - 1 else sampler_list[s + 1].fixed_frame | |
pelvis_new = points_gene[:, -fixed_frame, :9].cpu().numpy().reshape(cfg.batch_size, 3, 3) | |
trans_mats = np.repeat(np.eye(4)[np.newaxis, :, :], cfg.batch_size, axis=0) | |
for ip, pn in enumerate(pelvis_new): | |
_, ret_R, ret_t = rigid_transform_3D(np.matrix(pn), rest_pelvis, False) | |
ret_t[1] = 0.0 | |
rot_euler = R.from_matrix(ret_R).as_euler('zxy') | |
shift_euler = np.array([0, 0, rot_euler[2]]) | |
shift_rot_matrix2 = R.from_euler('zxy', shift_euler).as_matrix() | |
trans_mats[ip, :3, :3] = shift_rot_matrix2 | |
trans_mats[ip, :3, 3] = ret_t.reshape(-1) | |
mat = torch.from_numpy(trans_mats).to(device=cfg.device, dtype=torch.float32) | |
if fixed_frame > 0: | |
fixed_points = points_gene[:, -fixed_frame:] | |
if s == max_step - 1: | |
print('Saved Mat and Fixed Points', flush=True) | |
# np.save(os.path.join(cfg.exp_dir, f'{cfg.method_name}_mat.npy'), mat.cpu().numpy()) | |
# np.save(os.path.join(cfg.exp_dir, f'{cfg.method_name}_fixed_points.npy'), fixed_points.cpu().numpy()) | |
points_all = np.concatenate(points_all, axis=1) | |
points_all = points_all[:, cnt_seq_len - cnt_fixed_frame:] | |
return points_all | |
def sample_wrapper(trajectory, obj_locs): | |
trajectory = convert_trajectory(trajectory) | |
# obj_locs = {key: [data[key]['x'], data['key']['z']] for key in data.keys() if 'trajectory' not in key} | |
# cfg = compose(config_name="config_sample_synhsi") | |
with open('config/config_sample_synhsi.yaml') as f: | |
cfg = yaml.safe_load(f) | |
cfg = dotDict(cfg) | |
# @hydra.main(version_base=None, config_path="../config", config_name="config_sample_synhsi") | |
# def sample(cfg) -> None: | |
print(cfg) | |
# seed_everything(100) | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
# model_joints_to_smplx = init_model(cfg.model.model_smplx, device=device, eval=True) | |
model_joints_to_smplx = JointsToSMPLX(**cfg.model.model_smplx) | |
model_joints_to_smplx.load_state_dict(torch.load(cfg.model.model_smplx.ckpt)) | |
model_joints_to_smplx.to(device) | |
model_joints_to_smplx.eval() | |
# model_body = init_model(cfg.model.synhsi_body, device=device, eval=True) | |
model_body = Unet(**cfg.model.synhsi_body) | |
model_body.load_state_dict(torch.load(cfg.model.synhsi_body.ckpt)) | |
model_body.to(device) | |
model_body.eval() | |
# model_hand = init_model(cfg.model.synhsi_hand, device=device, eval=True) | |
# synhsi_dataset = hydra.utils.instantiate(cfg.dataset) | |
synhsi_dataset = TrumansDataset(**cfg.dataset) | |
sampler_body = hydra.utils.instantiate(cfg.sampler.pelvis) | |
# sampler_hand = hydra.utils.instantiate(cfg.sampler.right_hand) | |
sampler_body.set_dataset_and_model(synhsi_dataset, model_body) | |
# sampler_hand.set_dataset_and_model(None, model_hand) | |
samplers = {'body': sampler_body, 'hand': None} | |
# for scene_name in ['N3OpenArea']: | |
# trajectory = np.load(os.path.join(cfg.test_dir, cfg.exp_name, f'trajectories.npy'), allow_pickle=True).item() | |
# cfg.scene_name = scene_name | |
# cfg.action_type = trajectory['action_type'] | |
# if 'action_id' in trajectory.keys(): | |
# cfg.action_id = trajectory['action_id'] | |
# GP_LAYERS = ['GP_Layer'] | |
method_name = cfg.method_name | |
lid = 0 | |
base_speed = get_base_speed(cfg, trajectory, is_zup=False) | |
mat, goal_list, action_label_list, sampler_list = get_guidance(cfg, trajectory, samplers, act_type=cfg.action_type, | |
speed=int(0.6 * base_speed)) | |
points_all = sample_step(cfg, mat, obj_locs, goal_list, action_label_list, sampler_list) | |
# os.makedirs(cfg.exp_dir, exist_ok=True) | |
vertices = None | |
for i in range(cfg.batch_size): | |
keypoint_gene_torch = torch.from_numpy(points_all[i]).reshape(-1, cfg.dataset.nb_joints * 3).to(device) | |
pose, transl, left_hand, right_hand, vertices = joints_to_smpl(model_joints_to_smplx, keypoint_gene_torch, cfg.dataset.joints_ind, cfg.interp_s) | |
# output_data = {'transl': transl, 'body_pose': pose[:, 3:], 'global_orient': pose[:, :3], | |
# 'id': 0} | |
# print(output_data) | |
# with open(os.path.join(cfg.exp_dir, f'{method_name}_{lid}_{i}.pkl'), 'wb') as f: | |
# pkl.dump(output_data, f) | |
# vertices = np.load('/home/jiangnan/SyntheticHSI/Gradio_demo/vertices.npy', allow_pickle=True) | |
# np.save('/home/jiangnan/SyntheticHSI/Gradio_demo/vertices.npy', vertices) | |
return vertices.tolist() | |
# v = sample() | |
# | |
# | |
# return v | |
# def load_dataset_meta(cfg): | |
# metas = np.load(cfg.) | |
# if __name__ == '__main__': | |
# OmegaConf.register_resolver("times_three", times_three) | |
# OmegaConf.register_new_resolver("times", lambda x, y: int(x) * int(y)) | |
# sample() | |