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()