-
Notifications
You must be signed in to change notification settings - Fork 4
/
dataLoaderTalkSet.py
182 lines (167 loc) · 7.44 KB
/
dataLoaderTalkSet.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
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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
import os, torch, numpy, cv2, imageio, random, python_speech_features
import matplotlib.pyplot as plt
from scipy.io import wavfile
from glob import glob
from torchvision.transforms import RandomCrop
from scipy import signal
def get_noise_list(musanPath, rirPath):
augment_files = glob(os.path.join(musanPath, '*/*/*/*.wav'))
noiselist = {}
rir = numpy.load(rirPath)
for file in augment_files:
if not file.split('/')[-4] in noiselist:
noiselist[file.split('/')[-4]] = []
noiselist[file.split('/')[-4]].append(file)
return rir, noiselist
def augment_wav(audio, aug_type, rir, noiselist):
if aug_type == 'rir':
rir_gains = numpy.random.uniform(-7,3,1)
rir_filts = random.choice(rir)
rir = numpy.multiply(rir_filts, pow(10, 0.1 * rir_gains))
audio = signal.convolve(audio, rir, mode='full')[:len(audio)]
else:
noisecat = aug_type
noisefile = random.choice(noiselist[noisecat].copy())
snr = [random.uniform({'noise':[0,15],'music':[5,15]}[noisecat][0], {'noise':[0,15],'music':[5,15]}[noisecat][1])]
_, noiseaudio = wavfile.read(noisefile)
if len(noiseaudio) < len(audio):
shortage = len(audio) - len(noiseaudio)
noiseaudio = numpy.pad(noiseaudio, (0, shortage), 'wrap')
else:
noiseaudio = noiseaudio[:len(audio)]
noise_db = 10 * numpy.log10(numpy.mean(abs(noiseaudio ** 2)) + 1e-4)
clean_db = 10 * numpy.log10(numpy.mean(abs(audio ** 2)) + 1e-4)
noise = numpy.sqrt(10 ** ((clean_db - noise_db - snr) / 10)) * noiseaudio
audio = audio + noise
return audio.astype(numpy.int16)
def load_audio(data, data_path, length, start, end, audio_aug, rirlist = None, noiselist = None):
# Find the path of the audio data
data_type = data[0]
id_name = data[1][:8]
file_name = data[1].split('/')[0] + '_' + data[1].split('/')[1] + '_' + data[1].split('/')[2] + \
'_' + data[2].split('/')[0] + '_' + data[2].split('/')[1] + '_' + data[2].split('/')[2] + '.wav'
audio_file_path = os.path.join(data_path, data_type, id_name, file_name)
# Load audio, compute MFCC, cut it to the required length
_, audio = wavfile.read(audio_file_path)
if audio_aug == True:
augtype = random.randint(0,3)
if augtype == 1: # rir
audio = augment_wav(audio, 'rir', rirlist, noiselist)
elif augtype == 2:
audio = augment_wav(audio, 'noise', rirlist, noiselist)
elif augtype == 3:
audio = augment_wav(audio, 'music', rirlist, noiselist)
else:
audio = audio
feature = python_speech_features.mfcc(audio, 16000, numcep = 13, winlen = 0.025, winstep = 0.010)
length_audio = int(round(length * 100))
if feature.shape[0] < length_audio:
shortage = length_audio - feature.shape[0]
feature = numpy.pad(feature, ((0, shortage), (0,0)), 'wrap')
feature = feature[int(round(start * 100)):int(round(end * 100)),:]
return feature
def load_video(data, data_path, length, start, end, visual_aug):
# Find the path of the visual data
data_type = data[0]
id_name = data[1][:8]
file_name = data[1].split('/')[0] + '_' + data[1].split('/')[1] + '_' + data[1].split('/')[2] + \
'_' + data[2].split('/')[0] + '_' + data[2].split('/')[1] + '_' + data[2].split('/')[2] + '.mp4'
video_file_path = os.path.join(data_path, data_type, id_name, file_name)
# Load visual frame-by-frame, cut it to the required length
length_video = int(round((end - start) * 25))
video = cv2.VideoCapture(video_file_path)
faces = []
augtype = 'orig'
if visual_aug == True:
new = int(112*random.uniform(0.7, 1))
x, y = numpy.random.randint(0, 112 - new), numpy.random.randint(0, 112 - new)
M = cv2.getRotationMatrix2D((112/2,112/2), random.uniform(-15, 15), 1)
augtype = random.choice(['orig', 'flip', 'crop', 'rotate'])
num_frame = 0
while video.isOpened():
ret, frames = video.read()
if ret == True:
num_frame += 1
if num_frame >= int(round(start * 25)) and num_frame < int(round(end * 25)):
face = cv2.cvtColor(frames, cv2.COLOR_BGR2GRAY)
face = cv2.resize(face, (224,224))
face = face[int(112-(112/2)):int(112+(112/2)), int(112-(112/2)):int(112+(112/2))]
if augtype == 'orig':
faces.append(face)
elif augtype == 'flip':
faces.append(cv2.flip(face, 1))
elif augtype == 'crop':
faces.append(cv2.resize(face[y:y+new, x:x+new] , (112,112)))
elif augtype == 'rotate':
faces.append(cv2.warpAffine(face, M, (112,112)))
else:
break
video.release()
faces = numpy.array(faces)
if faces.shape[0] < length_video:
shortage = length_video - faces.shape[0]
faces = numpy.pad(faces, ((0,shortage), (0,0),(0,0)), 'wrap')
# faces = numpy.array(faces)[int(round(start * 25)):int(round(end * 25)),:,:]
return faces
def load_label(data, length, start, end):
labels_all = []
labels = []
data_type = data[0]
start_T, end_T, start_F, end_F = float(data[4]), float(data[5]), float(data[6]), float(data[7])
for i in range(int(round(length * 100))):
if data_type == 'TAudio':
labels_all.append(1)
elif data_type == 'FAudio' or data_type == 'FSilence':
labels_all.append(0)
else:
if i >= int(round(start_T * 100)) and i <= int(round(end_T * 100)):
labels_all.append(1)
else:
labels_all.append(0)
for i in range(int(round(length * 25))):
labels.append(int(round(sum(labels_all[i*4: (i+1)*4]) / 4)))
return labels[round(start*25): round(end*25)]
class loader_TalkSet(object):
def __init__(self, trial_file_name, data_path, audio_aug, visual_aug, musanPath, rirPath,**kwargs):
self.data_path = data_path
self.audio_aug = audio_aug
self.visual_aug = visual_aug
self.minibatch = []
self.rir, self.noiselist = get_noise_list(musanPath, rirPath)
mix_lst = open(trial_file_name).read().splitlines()
mix_lst = list(filter(lambda x: float(x.split()[3]) >= 1, mix_lst)) # filter the video less than 1s
# mix_lst = list(filter(lambda x: x.split()[0] == 'TSilence', mix_lst))
sorted_mix_lst = sorted(mix_lst, key=lambda data: (float(data.split()[3]), int(data.split()[-1])), reverse=True)
start = 0
while True:
length_total = float(sorted_mix_lst[start].split()[3])
batch_size = int(250 / length_total)
end = min(len(sorted_mix_lst), start + batch_size)
self.minibatch.append(sorted_mix_lst[start:end])
if end == len(sorted_mix_lst):
break
start = end
# self.minibatch = self.minibatch[0:5]
def __getitem__(self, index):
batch_lst = self.minibatch[index]
length_total = float(batch_lst[-1].split()[3])
length_total = (int(round(length_total * 100)) - int(round(length_total * 100)) % 4) / 100
audio_feature, video_feature, labels = [], [], []
duration = random.choice([1,2,4,6])
#duration = 6
length = min(length_total, duration)
if length == duration:
start = int(round(random.randint(0, round(length_total * 25) - round(length * 25)) * 0.04 * 100)) / 100
end = int(round((start + length) * 100)) / 100
else:
start, end = 0, length
for line in batch_lst:
data = line.split()
audio_feature.append(load_audio(data, self.data_path, length_total, start, end, audio_aug = self.audio_aug, rirlist = self.rir, noiselist = self.noiselist))
video_feature.append(load_video(data, self.data_path, length_total, start, end, visual_aug = self.visual_aug))
labels.append(load_label(data, length_total, start, end))
return torch.FloatTensor(numpy.array(audio_feature)), \
torch.FloatTensor(numpy.array(video_feature)), \
torch.LongTensor(numpy.array(labels))
def __len__(self):
return len(self.minibatch)