| import json |
| import torch |
| from torch.utils import data |
| import numpy as np |
| import librosa |
| import os |
| import sys |
|
|
| sys.path.append(os.path.join(os.path.dirname(__file__), "..")) |
| from emage_utils.motion_io import beat_format_load, MASK_DICT |
|
|
| class BEAT2Dataset(data.Dataset): |
| def __init__(self, cfg, split): |
| vid_meta = [] |
| for data_meta_path in cfg.data.meta_paths: |
| vid_meta.extend(json.load(open(data_meta_path, "r"))) |
| self.vid_meta = [item for item in vid_meta if item.get("mode") == split] |
| self.mean = 0 |
| self.std = 1 |
| self.joint_mask = MASK_DICT[cfg.model.joint_mask] if cfg.model.joint_mask is not None else None |
| self.data_list = self.vid_meta |
| self.fps = cfg.model.pose_fps |
| self.audio_sr = cfg.model.audio_sr |
|
|
| def __len__(self): |
| return len(self.data_list) |
| |
| @staticmethod |
| def normalize(motion, mean, std): |
| return (motion - mean) / (std + 1e-7) |
| |
| @staticmethod |
| def inverse_normalize(motion, mean, std): |
| return motion * std + mean |
|
|
| def __getitem__(self, item): |
| data_item = self.data_list[item] |
| smplx_data = beat_format_load(data_item["motion_path"], mask=self.joint_mask) |
| sdx, edx = data_item["start_idx"], data_item["end_idx"] |
| motion = smplx_data["poses"][sdx:edx] |
| SMPLX_FPS = 30 |
| downsample_factor = SMPLX_FPS // self.fps |
| motion = motion[::downsample_factor] |
| motion = self.normalize(motion, self.mean, self.std) |
| |
| audio, _ = librosa.load(data_item["audio_path"], sr=self.audio_sr) |
| sdx_audio = sdx * int((1 / SMPLX_FPS) * self.audio_sr) |
| edx_audio = edx * int((1 / SMPLX_FPS) * self.audio_sr) |
| audio = audio[sdx_audio:edx_audio] |
| |
| motion_tensor = torch.from_numpy(motion).float() |
| audio_tensor = torch.from_numpy(audio).float() |
| |
| return dict( |
| motion=motion_tensor, |
| audio=audio_tensor, |
| ) |
|
|
| class BEAT2DatasetEamge(BEAT2Dataset): |
| def __init__(self, cfg, split): |
| super().__init__(cfg, split) |
|
|
| def __getitem__(self, item): |
| data_item = self.data_list[item] |
| smplx_data = beat_format_load(data_item["motion_path"], mask=None) |
| sdx, edx = data_item["start_idx"], data_item["end_idx"] |
| motion = smplx_data["poses"][sdx:edx] |
| expressions = smplx_data["expressions"][sdx:edx] |
| trans = smplx_data["trans"][sdx:edx] |
| SMPLX_FPS = 30 |
| downsample_factor = SMPLX_FPS // self.fps |
| motion = motion[::downsample_factor] |
| motion = self.normalize(motion, self.mean, self.std) |
| |
| audio, _ = librosa.load(data_item["audio_path"], sr=self.audio_sr) |
| sdx_audio = sdx * int((1 / SMPLX_FPS) * self.audio_sr) |
| edx_audio = edx * int((1 / SMPLX_FPS) * self.audio_sr) |
| audio = audio[sdx_audio:edx_audio] |
| |
| motion_tensor = torch.from_numpy(motion).float() |
| audio_tensor = torch.from_numpy(audio).float() |
| expressions_tesnor = torch.from_numpy(expressions).float() |
| trans_tensor = torch.from_numpy(trans).float() |
|
|
| return dict( |
| motion=motion_tensor, |
| audio=audio_tensor, |
| expressions=expressions_tesnor, |
| trans=trans_tensor, |
| ) |
|
|
|
|
| class BEAT2DatasetEamgeFootContact(BEAT2Dataset): |
| def __init__(self, cfg, split): |
| super().__init__(cfg, split) |
|
|
| def __getitem__(self, item): |
| data_item = self.data_list[item] |
| smplx_data = beat_format_load(data_item["motion_path"], mask=None) |
| sdx, edx = data_item["start_idx"], data_item["end_idx"] |
| motion = smplx_data["poses"][sdx:edx] |
| expressions = smplx_data["expressions"][sdx:edx] |
| trans = smplx_data["trans"][sdx:edx] |
| foot_contact = np.load(data_item["motion_path"].replace("smplxflame_30", "footcontact").replace(".npz", ".npy"))[sdx:edx] |
|
|
| SMPLX_FPS = 30 |
| downsample_factor = SMPLX_FPS // self.fps |
| motion = motion[::downsample_factor] |
| motion = self.normalize(motion, self.mean, self.std) |
| |
| audio, _ = librosa.load(data_item["audio_path"], sr=self.audio_sr) |
| sdx_audio = sdx * int((1 / SMPLX_FPS) * self.audio_sr) |
| edx_audio = edx * int((1 / SMPLX_FPS) * self.audio_sr) |
| audio = audio[sdx_audio:edx_audio] |
| |
| motion_tensor = torch.from_numpy(motion).float() |
| audio_tensor = torch.from_numpy(audio).float() |
| expressions_tesnor = torch.from_numpy(expressions).float() |
| trans_tensor = torch.from_numpy(trans).float() |
| foot_contact_tensor = torch.from_numpy(foot_contact).float() |
| |
|
|
| return dict( |
| motion=motion_tensor, |
| audio=audio_tensor, |
| expressions=expressions_tesnor, |
| trans=trans_tensor, |
| foot_contact=foot_contact_tensor, |
| ) |
|
|
