|
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):
|
|
|
|
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)
|
|
|
|
_, audio = wavfile.read(audio_file_path)
|
|
|
|
if audio_aug == True:
|
|
augtype = random.randint(0,3)
|
|
if augtype == 1:
|
|
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):
|
|
|
|
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)
|
|
|
|
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')
|
|
|
|
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))
|
|
|
|
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
|
|
|
|
|
|
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])
|
|
|
|
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) |