| from torch.utils.data import Dataset
|
| import os
|
| import torch
|
| import random
|
| import numpy as np
|
|
|
| class ReactDataset(Dataset):
|
| def __init__(
|
| self,
|
| root_dir,
|
| split="train",
|
| clip_length=750,
|
| random_crop=True
|
| ):
|
| self.clip_length = clip_length
|
| self.random_crop = random_crop
|
| base = os.path.join(root_dir, split)
|
| self.audio_dir = os.path.join(base, "audio-features")
|
| self.dmm_dir = os.path.join(base, "coefficients")
|
| self.emo_dir = os.path.join(base, "facial-attributes")
|
|
|
| speaker_emo_root = os.path.join(self.emo_dir, "speaker")
|
| self.items = []
|
| for root, _, files in os.walk(speaker_emo_root):
|
| session = os.path.basename(root)
|
| for f in files:
|
| if not f.lower().endswith(".npy"):
|
| continue
|
| self.items.append(os.path.join(session, f))
|
| self.items.sort()
|
| assert len(self.items) > 0, "No .npy files found in " + speaker_emo_root
|
|
|
| mean_face_path = os.path.join("external_tools", "mean_face.npy")
|
| std_face_path = os.path.join("external_tools", "std_face.npy")
|
|
|
| self.mean_face = np.load(mean_face_path).astype(np.float32)
|
| self.std_face = np.load(std_face_path).astype(np.float32)
|
|
|
| self.mean_face = torch.from_numpy(self.mean_face).unsqueeze(0)
|
| self.std_face = torch.from_numpy(self.std_face).unsqueeze(0)
|
|
|
| def __len__(self):
|
| return len(self.items)
|
|
|
| def load_data(self, modality_dir, speaker_listener, current, target_dim=None):
|
| path = os.path.join(modality_dir, speaker_listener, current)
|
| t = np.load(path).astype(np.float32)
|
| t = torch.from_numpy(t).squeeze()
|
|
|
| if target_dim is not None:
|
| assert t.dim() == target_dim
|
|
|
| return t
|
|
|
| def __getitem__(self, index):
|
| current = self.items[index]
|
|
|
| sp_audio = self.load_data(self.audio_dir, "speaker", current, target_dim=2)
|
| sp_dmm = self.load_data(self.dmm_dir, "speaker", current, target_dim=2)
|
| sp_emo = self.load_data(self.emo_dir, "speaker", current, target_dim=2)
|
| li_emo = self.load_data(self.emo_dir, "listener", current, target_dim=2)
|
|
|
| T = min(sp_audio.size(0), sp_dmm.size(0), sp_emo.size(0), li_emo.size(0))
|
| sp_audio = sp_audio[:T]
|
| sp_dmm = sp_dmm[:T]
|
| sp_emo = sp_emo[:T]
|
| li_emo = li_emo[:T]
|
|
|
| sp_dmm = (sp_dmm - self.mean_face) / self.std_face
|
|
|
| if T > self.clip_length:
|
| if self.random_crop:
|
| start = random.randint(0, T - self.clip_length)
|
| else:
|
| start = (T - self.clip_length) // 2
|
| sp_audio = sp_audio[start: start + self.clip_length]
|
| sp_dmm = sp_dmm[start: start + self.clip_length]
|
| sp_emo = sp_emo[start: start + self.clip_length]
|
| li_emo = li_emo[start: start + self.clip_length]
|
| motion_length = self.clip_length
|
| else:
|
| pad = self.clip_length - T
|
| sp_audio = torch.cat([sp_audio, torch.zeros(pad, 768)], dim=0)
|
| sp_dmm = torch.cat([sp_dmm, torch.zeros(pad, 58)], dim=0)
|
| sp_emo = torch.cat([sp_emo, torch.zeros(pad, 25)], dim=0)
|
| li_emo = torch.cat([li_emo, torch.zeros(pad, 25)], dim=0)
|
| motion_length = T
|
|
|
| assert sp_audio.size() == (750, 768)
|
| assert sp_dmm.size() == (750, 58)
|
| assert sp_emo.size() == (750, 25)
|
| assert li_emo.size() == (750, 25)
|
|
|
| return {
|
| "speaker_audio": sp_audio,
|
| "speaker_3dmm": sp_dmm,
|
| "speaker_emotion": sp_emo,
|
| "listener_emotion": li_emo,
|
| "motion_length": torch.tensor(motion_length, dtype=torch.long)
|
| } |
