import numpy as np import random, torch, joblib from torch.utils.data import Dataset from mmengine.registry import DATASETS, FUNCTIONS import time, os from torch.nn.utils.rnn import pad_sequence from ..utils.motion import recover_from_local_position_fast FPS = 30 Audio_FPS = 75 @DATASETS.register_module() class AnyMoDataset(Dataset): ''' data format: dict( motion_272: np.ndarray (T, 272), captions: [str] ) ''' def __init__(self, data_path_file='', unit_length=4, mean_path='', std_path='', audio_unit_length=8, min_motion_length=60, traj_mean_path='', traj_std_path='', max_motion_length=300, interval=1): super().__init__() self.min_motion_length = min_motion_length self.max_motion_length = max_motion_length self.unit_length = unit_length self.audio_unit_length = audio_unit_length if isinstance(data_path_file, list): self.data_paths = [] for path in data_path_file: with open(path, 'r') as f: self.data_paths.extend([line.strip() for line in f.readlines()]) elif isinstance(data_path_file, str): with open(data_path_file, 'r') as f: self.data_paths = [line.strip() for line in f.readlines()] else: raise NotImplementedError self.data_paths = self.data_paths[::interval] self.motion_mean = torch.from_numpy(np.load(mean_path).astype(np.float32)) self.motion_std = torch.from_numpy(np.load(std_path).astype(np.float32)) self.traj_mean = torch.from_numpy(np.load(traj_mean_path).astype(np.float32)) self.traj_std = torch.from_numpy(np.load(traj_std_path).astype(np.float32)) def __len__(self): return len(self.data_paths) def __getitem__(self, index): data_infos = joblib.load(self.data_paths[index]) data = [] for data_info in data_infos: # 1. load motion motion = torch.as_tensor(data_info['motion_272'], dtype=torch.float32) caption = random.choice(data_info['captions']) audio_feat = torch.as_tensor(data_info['audio_feat'], dtype=torch.float32) # traj_data = torch.as_tensor(data_info['root_traj'], dtype=torch.float32) traj_data = recover_from_local_position_fast(motion[None]).squeeze(0)[:, 0] # import ipdb; ipdb.set_trace() # 有一些比较长的音频，这部分不用caption标注 if audio_feat.shape[0] // Audio_FPS > 15: caption = '' # 2. clip motion T = motion.shape[0] motion_length = (T // self.unit_length) * self.unit_length if motion_length > self.max_motion_length: motion_length = self.max_motion_length idx = random.randint(0, T - motion_length) motion = motion[idx:idx+motion_length] traj_data = traj_data[idx:idx+motion_length] traj_data = (traj_data - self.traj_mean) / self.traj_std motion = (motion - self.motion_mean) / self.motion_std pad_motion = torch.zeros(self.max_motion_length - motion_length, motion.shape[1]).to(motion) motion = torch.cat([motion, pad_motion], dim=0) pad_traj = torch.zeros(self.max_motion_length - motion_length, traj_data.shape[1]).to(traj_data) traj_data = torch.cat([traj_data, pad_traj], dim=0) # 3. clip audio audio_length = motion_length / FPS * Audio_FPS audio_length = int((audio_length // self.audio_unit_length) * self.audio_unit_length) audio_idx = int(idx / FPS * Audio_FPS) audio_feat = audio_feat[audio_idx:audio_idx+audio_length] data.append(dict( motion=motion, traj_data=traj_data, text_data=caption, audio_data=audio_feat, motion_length=motion_length, # audio_length=audio_length, audio_length=audio_feat.shape[0], mean=self.motion_mean, std=self.motion_std, )) return data @DATASETS.register_module() class AnyMoDataset_v1(AnyMoDataset): # 允许数据缺失 def __getitem__(self, index): data_infos = joblib.load(self.data_paths[index]) data = [] for data_info in data_infos: # 1. load motion motion = torch.as_tensor(data_info['motion_272'], dtype=torch.float32) m_path = data_info['motion_272_path'] caption = random.choice(data_info['captions']) has_audio = 'audio_feat' in data_info traj_data = recover_from_local_position_fast(motion[None]).squeeze(0)[:, 0] # 2. clip motion T = motion.shape[0] motion_length = (T // self.unit_length) * self.unit_length if motion_length > self.max_motion_length: motion_length = self.max_motion_length idx = random.randint(0, T - motion_length) motion = motion[idx:idx+motion_length] traj_data = traj_data[idx:idx+motion_length] traj_data = (traj_data - self.traj_mean) / self.traj_std motion = (motion - self.motion_mean) / self.motion_std pad_motion = torch.zeros(self.max_motion_length - motion_length, motion.shape[1]).to(motion) motion = torch.cat([motion, pad_motion], dim=0) pad_traj = torch.zeros(self.max_motion_length - motion_length, traj_data.shape[1]).to(traj_data) traj_data = torch.cat([traj_data, pad_traj], dim=0) # 3. clip audio if has_audio: # 有一些比较长的音频，这部分不用caption标注 audio_feat = torch.as_tensor(data_info['audio_feat'], dtype=torch.float32) if audio_feat.shape[0] // Audio_FPS > 15: caption = '' audio_length = motion_length / FPS * Audio_FPS audio_length = int((audio_length // self.audio_unit_length) * self.audio_unit_length) audio_idx = int(idx / FPS * Audio_FPS) audio_feat = audio_feat[audio_idx:audio_idx+audio_length] beat_path = os.path.join("data/omnihumo/motion_caption_audio_beat", data_info['motion_272_path']) # print(np.load(beat_path).shape, motion_length) audio_beat = torch.from_numpy(np.load(beat_path)[0, :motion_length]) else: audio_feat = torch.zeros((0, 512)).float() audio_length = 0 audio_beat = torch.zeros((0, )).float() data.append(dict( motion=motion, m_path=m_path, traj_data=traj_data, text_data=caption, audio_data=audio_feat, motion_length=motion_length, audio_length=audio_feat.shape[0], audio_beat=audio_beat, mean=self.motion_mean, std=self.motion_std, )) return data @DATASETS.register_module() class AnyMoDataset_v2(Dataset): # 对 text 采样 + 允许数据缺失 def __init__(self, audio_path_file='', caption_path_file='', unit_length=4, mean_path='', std_path='', audio_unit_length=8, min_motion_length=60, caption_sample_ratio=0.3, traj_mean_path='', traj_std_path='', max_motion_length=300, interval=1): super().__init__() self.min_motion_length = min_motion_length self.max_motion_length = max_motion_length self.unit_length = unit_length self.audio_unit_length = audio_unit_length audio_paths = [] for path in audio_path_file: with open(path, 'r') as f: audio_paths.extend([line.strip() for line in f.readlines()]) with open(caption_path_file, 'r') as f: caption_paths = [line.strip() for line in f.readlines()] self.audio_paths = audio_paths[::interval] self.caption_paths = caption_paths[::interval] self.caption_sample_ratio = caption_sample_ratio # 对caption采样率 self.motion_mean = torch.from_numpy(np.load(mean_path).astype(np.float32)) self.motion_std = torch.from_numpy(np.load(std_path).astype(np.float32)) self.traj_mean = torch.from_numpy(np.load(traj_mean_path).astype(np.float32)) self.traj_std = torch.from_numpy(np.load(traj_std_path).astype(np.float32)) def __len__(self): return len(self.audio_paths) + int(len(self.caption_paths) * self.caption_sample_ratio) def __getitem__(self, index): if index < len(self.audio_paths): data_infos = joblib.load(self.audio_paths[index]) else: index = index - len(self.audio_paths) cap_size = len(self.caption_paths) size = int(cap_size * self.caption_sample_ratio) begin = random.randint(0, cap_size - size) data_infos = joblib.load(self.caption_paths[begin + index]) data = [] for data_info in data_infos: # 1. load motion motion = torch.as_tensor(data_info['motion_272'], dtype=torch.float32) caption = random.choice(data_info['captions']) has_audio = 'audio_feat' in data_info traj_data = recover_from_local_position_fast(motion[None]).squeeze(0)[:, 0] # 2. clip motion T = motion.shape[0] motion_length = (T // self.unit_length) * self.unit_length if motion_length > self.max_motion_length: motion_length = self.max_motion_length idx = random.randint(0, T - motion_length) motion = motion[idx:idx+motion_length] traj_data = traj_data[idx:idx+motion_length] traj_data = (traj_data - self.traj_mean) / self.traj_std motion = (motion - self.motion_mean) / self.motion_std pad_motion = torch.zeros(self.max_motion_length - motion_length, motion.shape[1]).to(motion) motion = torch.cat([motion, pad_motion], dim=0) pad_traj = torch.zeros(self.max_motion_length - motion_length, traj_data.shape[1]).to(traj_data) traj_data = torch.cat([traj_data, pad_traj], dim=0) # 3. clip audio if has_audio: # 有一些比较长的音频，这部分不用caption标注 audio_feat = torch.as_tensor(data_info['audio_feat'], dtype=torch.float32) if audio_feat.shape[0] // Audio_FPS > 15: caption = '' audio_length = motion_length / FPS * Audio_FPS audio_length = int((audio_length // self.audio_unit_length) * self.audio_unit_length) audio_idx = int(idx / FPS * Audio_FPS) audio_feat = audio_feat[audio_idx:audio_idx+audio_length] # import ipdb; ipdb.set_trace() else: audio_feat = torch.zeros((0, 512)).float() audio_length = 0 data.append(dict( motion=motion, traj_data=traj_data, text_data=caption, audio_data=audio_feat, motion_length=motion_length, audio_length=audio_feat.shape[0], mean=self.motion_mean, std=self.motion_std, )) return data @FUNCTIONS.register_module() def anymo_collate_fn(batch): new_batch = [] for item in batch: new_batch.extend(item) motion = [item['motion'] for item in new_batch] # [motion: T, C] motion = torch.stack(motion) traj_data = [item['traj_data'] for item in new_batch] traj_data = torch.stack(traj_data) motion_length = torch.as_tensor([item['motion_length'] for item in new_batch]) audio_length = torch.as_tensor([item['audio_length'] for item in new_batch]) text_data = [item['text_data'] for item in new_batch] motion = motion[:, :motion_length.max()] traj_data = traj_data[:, :motion_length.max()] mean = new_batch[0]['mean'] std = new_batch[0]['std'] audio_data = [item['audio_data'] for item in new_batch] audio_data = pad_sequence(audio_data, batch_first=True) if 'audio_beat' in new_batch[0]: audio_beat = [item['audio_beat'] for item in new_batch] audio_beat = pad_sequence(audio_beat, batch_first=True) else: audio_beat = None # print(motion.shape, traj_data.shape, audio_data.shape, len(text_data)) return dict( motion=motion, traj_data=traj_data, audio_data=audio_data, motion_length=motion_length, audio_length=audio_length, text_data=text_data, audio_beat=audio_beat, mean=mean, std=std, ) @FUNCTIONS.register_module() def anymo_collate_fn_v1(batch): new_batch = [] for item in batch: new_batch.extend(item) motion = [item['motion'] for item in new_batch] # [motion: T, C] m_path = [item['m_path'] for item in new_batch] # [motion: T, C] motion = torch.stack(motion) traj_data = [item['traj_data'] for item in new_batch] traj_data = torch.stack(traj_data) motion_length = torch.as_tensor([item['motion_length'] for item in new_batch]) audio_length = torch.as_tensor([item['audio_length'] for item in new_batch]) text_data = [item['text_data'] for item in new_batch] motion = motion[:, :motion_length.max()] traj_data = traj_data[:, :motion_length.max()] mean = new_batch[0]['mean'] std = new_batch[0]['std'] audio_data = [item['audio_data'] for item in new_batch] audio_data = pad_sequence(audio_data, batch_first=True) if 'audio_beat' in new_batch[0]: audio_beat = [item['audio_beat'] for item in new_batch] audio_beat = pad_sequence(audio_beat, batch_first=True) else: audio_beat = None # print(motion.shape, traj_data.shape, audio_data.shape, len(text_data)) return dict( motion=motion, traj_data=traj_data, audio_data=audio_data, motion_length=motion_length, audio_length=audio_length, text_data=text_data, audio_beat=audio_beat, m_path=m_path, mean=mean, std=std, )