-
Notifications
You must be signed in to change notification settings - Fork 2
/
EmodbDataset.py
65 lines (51 loc) · 2.01 KB
/
EmodbDataset.py
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
import os
import torch
import torchaudio
import pandas as pd
import numpy as np
class EmodbDataset(object):
"""
Create a Dataset for Emodb. Each item is a tuple of the form:
(waveform, sample_rate, emotion)
"""
_ext_audio = '.wav'
_emotions = { 'W': 1, 'L': 2, 'E': 3, 'A': 4, 'F': 5, 'T': 6, 'N': 7 } # W = anger, L = boredom, E = disgust, A = anxiety/fear, F = happiness, T = sadness, N = neutral
def __init__(self, root='download'):
"""
Args:
root (string): Directory containing the wav folder
"""
self.root = root
# Iterate through all audio files
data = []
for _, _, files in os.walk(root):
for file in files:
if file.endswith(self._ext_audio):
# Construct file identifiers
identifiers = [file[0:2], file[2:5], file[5], file[6], os.path.join('wav', file)]
# Append identifier to data
data.append(identifiers)
# Create pandas dataframe
self.df = pd.DataFrame(data, columns=['speaker_id', 'code', 'emotion', 'version', 'file'], dtype=np.float32)
# Map emotion labels to numeric values
self.df['emotion'] = self.df['emotion'].map(self._emotions).astype(np.float32)
def __len__(self):
return len(self.df)
def __getitem__(self, idx):
if torch.is_tensor(idx):
idx = idx.tolist()
audio_name = os.path.join(self.root, self.df.loc[idx, 'file'])
waveform, sample_rate = torchaudio.load(audio_name)
emotion = self.df.loc[idx, 'emotion']
sample = {
'waveform': waveform,
'sample_rate': sample_rate,
'emotion': emotion
}
return sample
# Example: Load Emodb dataset
# emodb_dataset = EmodbDataset('/home/alanwuha/Documents/Projects/datasets/emodb/download')
# Example: Iterate through samples
# for i in range(len(emodb_dataset)):
# sample = emodb_dataset[i]
# print(i, sample)