File size: 3,938 Bytes
47cac2b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | 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] #e.g "session0/123.npy"
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)
} |