import warnings, math from os.path import join as pjoin from pathlib import Path import numpy as np import torch from tqdm import tqdm from typing import Literal from torch.utils.data import Dataset from .. import dist as dist_utils from .tensors import motion_action_collate, motion_text_collate, motion_text_treble_collate from ..metric.interaction import give_contact_label from ...feature.single_motioncode import MotionCoder from ...constant import INTRA_TIP_CONTACT_THRESH, TIP_PALM_CONTACT_THRESH, PALM_PALM_CONTACT_THRESH, SKELETON_CHAIN class HandXDataset(Dataset): contact_label: bool normalize: bool repr: Literal['joint_pos', 'joint_rot', 'joint_pos_w_scalar_rot', "joint_pos_w_axisangle_rot"] data_file_name: str data_dir: str def __init__(self, split: str, debug=False, *args, **kwargs): super(HandXDataset, self).__init__() self.debug = debug self.data_dir = kwargs['data_dir'] self.data_file_name = kwargs['data_file_name'] self.repr = kwargs['repr'] self.normalize = kwargs['normalize'] self.contact_label = kwargs.get("contact_label", False) self.ratio = kwargs.get("ratio", 1.0) if split == 'train': self.data_file_name = "train_" + self.data_file_name self.mano_file_name = 'train_mano.npz' elif split == 'val': self.data_file_name = 'test_' + self.data_file_name self.mano_file_name = 'test_mano.npz' self.split = split self.load_motion() if self.normalize: self._calc_mean_std() self.collate_fn = motion_text_treble_collate def _calc_mean_std(self): if self.split == 'val': dist_utils.barrier() print("Val split doesn't need create mean and std files.") mean_std_path = pjoin(self.data_dir, f'mean_std_{self.repr}') assert Path(mean_std_path).exists(), f"Mean and std not found at {mean_std_path}. Please run the dataset preparation script." self.mean = np.load(pjoin(mean_std_path, 'mean.npy')) self.std = np.load(pjoin(mean_std_path, 'std.npy')) return total_frames = 0 first_motion = next(iter(self.data_dict.values()))['motion'] feature_shape = first_motion.shape[1:] sum_of_data = np.zeros(feature_shape, dtype=np.float64) sum_of_squares = np.zeros(feature_shape, dtype=np.float64) for data in self.data_dict.values(): motion_data = data['motion'] if motion_data.ndim != 3: raise ValueError(f"Expected motion data to be 3D, but got {motion_data.ndim}D shape.") num_frames_in_batch = motion_data.shape[0] total_frames += num_frames_in_batch sum_of_data += np.sum(motion_data, axis=0) sum_of_squares += np.sum(np.square(motion_data), axis=0) if total_frames == 0: warnings.warn("No frames found in the dataset. Mean and std will be zero.") self.mean = np.zeros(feature_shape) self.std = np.zeros(feature_shape) else: self.mean = sum_of_data / total_frames variance = (sum_of_squares / total_frames) - np.square(self.mean) variance[variance < 0] = 0 self.std = np.sqrt(variance) self.std[self.std < 1e-4] = 1.0 if dist_utils.is_main_process(): save_path = (Path(self.data_dir) / f'mean_std_{self.repr}').as_posix() Path(save_path).mkdir(parents=True, exist_ok=True) # mylogger.info(f"Saving mean and std to {save_path}") np.save((Path(save_path) / 'mean.npy').as_posix(), self.mean) np.save((Path(save_path) / 'std.npy').as_posix(), self.std) def inv_transform(self, data): if isinstance(data, torch.Tensor): tmp_mean = torch.from_numpy(self.mean).to(data.device).float() tmp_std = torch.from_numpy(self.std).to(data.device).float() ret = data * tmp_std.reshape(-1) + tmp_mean.reshape(-1) # [B, T, J*C] return ret elif isinstance(data, np.ndarray): ret = data * self.std.reshape(-1) + self.mean.reshape(-1) # [B, T, J*C] return ret else: raise TypeError(f"Unsupported data type: {type(data)}. Expected torch.Tensor or np.ndarray.") def _get_axisangle_rotation(self, mano_pose:np.ndarray) -> np.ndarray: ''' mano_pose: (T, 48) return: (T, J, 3) ''' mano_pose = mano_pose.reshape(mano_pose.shape[0], -1, 3) # (T, 16, 3) zero_padding = np.zeros((mano_pose.shape[0], 21-16, 3)) # (T, 5, 3) return np.concatenate([mano_pose, zero_padding], axis=1) # (T, 21, 3) def _get_scalar_rotation(self, single_motion_seq:np.ndarray, side:Literal['left', 'right']): ''' single_motion_seq: (T, J, 3) ''' temp_motioncoder = MotionCoder(single_motion_seq, isright=(side=='right')) temp_motioncoder.get_local_coordinate() local_motion = temp_motioncoder.local_motion # (T, J, 3) additional_scalar_rotation = np.zeros((single_motion_seq.shape[0], single_motion_seq.shape[1])) # (T, J) for skeleton_chain in SKELETON_CHAIN: additional_scalar_rotation[:, skeleton_chain[0]] = 0 additional_scalar_rotation[:, skeleton_chain[-1]] = 0 for s in range(1, len(skeleton_chain) - 1): j = skeleton_chain[s] pre_j = skeleton_chain[s - 1] nxt_j = skeleton_chain[s + 1] v1 = (local_motion[:, j, :] - local_motion[:, pre_j, :])[:, [0, 2]] # (T, 2) v2 = (local_motion[:, nxt_j, :] - local_motion[:, j, :])[:, [0, 2]] # (T, 2) v1_direction_angle = np.arctan2(v1[:, 1], v1[:, 0]) # (T,) v2_direction_angle = np.arctan2(v2[:, 1], v2[:, 0]) # (T,) angle_diff = v2_direction_angle - v1_direction_angle # (T,) if side == 'right': angle_diff = -angle_diff additional_scalar_rotation[:, j] = angle_diff return additional_scalar_rotation # (T, J) def load_motion(self): motion_data = dict(np.load(pjoin(self.data_dir, self.data_file_name), allow_pickle=True)) for key in motion_data: motion_data[key] = motion_data[key].item() # mano_data = dict(np.load(pjoin(self.data_dir, self.mano_file_name), allow_pickle=True)) # for key in tqdm(mano_data, desc=f"RANK {dist_utils.get_rank()} | loading mano data for {self.split} split"): # mano_data[key] = mano_data[key].item() self.data_dict = dict() self.name_list = [] self.length_list = [] for clip_name in tqdm(sorted(motion_data.keys()), desc=f"RANK {dist_utils.get_rank()} | processing {self.split} data"): motion = motion_data[clip_name]['motion'] # (T, 2, J, 3) if self.repr == 'joint_pos_w_axisangle_rot': left_mano_pose = mano_data[clip_name]['left_pose'] # (T, 48) right_mano_pose = mano_data[clip_name]['right_pose'] # (T, 48) left_rot = self._get_axisangle_rotation(left_mano_pose) # (T, J, 3) right_rot = self._get_axisangle_rotation(right_mano_pose) # (T, J, 3) motion = np.concatenate([ motion, np.stack([left_rot, right_rot], axis=1) # (T, 2, J, 3) ], axis=-1) # (T, 2, J, 6) motion = motion.reshape(motion.shape[0], -1, 6) # (T, 2J, 6) transl = np.mean((motion[:, 0, :3] + motion[:, 21, :3]) / 2, axis=0) # (3,) motion[:, :, :3] -= transl elif self.repr == 'joint_pos_w_scalar_rot': left_rot_scalar = self._get_scalar_rotation(motion[:, 0], side='left') # (T, J) left_rot_scalar = np.nan_to_num(left_rot_scalar, nan=0.0, posinf=0.0, neginf=0.0) right_rot_scalar = self._get_scalar_rotation(motion[:, 1], side='right') # (T, J) right_rot_scalar = np.nan_to_num(right_rot_scalar, nan=0.0, posinf=0.0, neginf=0.0) motion = np.concatenate([ motion, np.stack([left_rot_scalar, right_rot_scalar], axis=1)[:, :, :, np.newaxis] # (T, 2, J, 1) ], axis=-1) # (T, 2, J, 4) motion = motion.reshape(motion.shape[0], -1, 4) # (T, 2J, 4) transl = np.mean((motion[:, 0, :3] + motion[:, 21, :3]) / 2, axis=0) # (3,) motion[:, :, :3] -= transl elif self.repr == 'joint_pos': transl = np.mean((motion[:, 0] + motion[:, 21]) / 2, axis=0) # (3,) motion[:, :, :3] -= transl pass else: raise NotImplementedError(f"Representation {self.repr} not implemented.") assert len(motion_data[clip_name]['left_annotation']) == len(motion_data[clip_name]['right_annotation']) == len(motion_data[clip_name]['interaction_annotation']), f"Annotation length mismatch for clip {clip_name}" annotation_count = len(motion_data[clip_name]['left_annotation']) for j in range(annotation_count): name = f"{clip_name}_ann{j}" self.name_list.append(name) self.length_list.append(motion.shape[0]) self.data_dict[name] = dict( motion=motion, text={ 'left': motion_data[clip_name]['left_annotation'][j], 'right': motion_data[clip_name]['right_annotation'][j], 'two_hands_relation': motion_data[clip_name]['interaction_annotation'][j] } ) self.length_list = np.array(self.length_list) if self.split == 'train' and self.ratio < 1.0: samples_count = int(math.ceil(len(self.name_list) * self.ratio)) random_indices = np.random.choice(len(self.name_list), size=samples_count, replace=False) self.name_list = [self.name_list[i] for i in random_indices] self.length_list = self.length_list[random_indices] self.data_dict = {name: self.data_dict[name] for name in self.name_list} def __getitem__(self, index): if self.debug: print(f"motion name: {self.name_list[index]}") print(f"self.name_list is UNIQUE: {len(self.name_list) == len(set(self.name_list))}") motion = self.data_dict[self.name_list[index]]['motion'] # (T, 2J, C) m_length = self.length_list[index] text = self.data_dict[self.name_list[index]]['text'] if self.contact_label == True: contact_label = give_contact_label( motion[:, :, :3].reshape(motion.shape[0], 2, -1, 3), # (T, 2, J, 3) tip_tip_threshold=INTRA_TIP_CONTACT_THRESH, tip_palm_threshold=TIP_PALM_CONTACT_THRESH, palm_palm_threshold=PALM_PALM_CONTACT_THRESH ) if self.normalize: motion = (motion - self.mean[np.newaxis]) / self.std if self.contact_label: return motion.transpose((1, 2, 0)), m_length, text, contact_label else: return motion.transpose((1, 2, 0)), m_length, text def __len__(self): return len(self.data_dict)