File size: 5,933 Bytes
44b657f bd21b55 a350e21 bd21b55 a350e21 44b657f bd21b55 44b657f bd21b55 44b657f bd21b55 44b657f bd21b55 a350e21 bd21b55 a350e21 44b657f bd21b55 44b657f a350e21 44b657f a350e21 bd21b55 44b657f a350e21 44b657f bd21b55 44b657f a350e21 44b657f | 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 | #!/usr/bin/env python3
"""Generates precomputed benchmark audio examples for the TEC HF Space demo."""
import json
import os
import sys
import numpy as np
import tensorflow.compat.v1 as tf
REPO_ROOT = '/usr/local/google/home/quanw/Code/github/tec'
if REPO_ROOT not in sys.path:
sys.path.insert(0, REPO_ROOT)
from tec import data_prep # noqa: E402
from tec import inference # noqa: E402
TFRECORD_PATH = (
'/usr/local/google/home/quanw/Data/tec_datasets/tfrecords/'
'single_test_clean.tfrecord'
)
CHECKPOINT_ROOT = (
'/usr/local/google/home/quanw/Data/tec_datasets/checkpoints'
)
EXAMPLES_DIR = os.path.join(
os.path.dirname(os.path.abspath(__file__)), 'examples'
)
TARGET_SAMPLES = [
{
'id': 1,
'utt_id': '4446_2275_000048_000001__LJ012-0176',
'lj_id': 'LJ012-0176',
'clean_ref': '"We\'ve tortured each other enough for tonight.',
'interfering_text': (
'To obtain possession of this, Thurtell with his two associates'
' resolved to kill him.'
),
},
{
'id': 2,
'utt_id': '1995_1837_000024_000001__LJ030-0183',
'lj_id': 'LJ030-0183',
'clean_ref': 'The chair was empty; but he knew.',
'interfering_text': (
"As he told the driver, quote, Let's get out of here; we are hit,"
' end quote,'
),
},
{
'id': 3,
'utt_id': '260_123286_000017_000001__LJ040-0208',
'lj_id': 'LJ040-0208',
'clean_ref': 'Have we followed that road?',
'interfering_text': (
'Carro reported that when questioned about his mother Lee said,'
' quote,'
),
},
{
'id': 4,
'utt_id': '8224_274384_000017_000000__LJ043-0134',
'lj_id': 'LJ043-0134',
'clean_ref': (
'The parliament and the Scots laid their proposals before the king.'
),
'interfering_text': (
'The Commission has concluded that on April ten, nineteen'
' sixty-three,'
),
},
]
MODELS = [
('tec_single', 'TecSingleInterfering'),
('tec_multi', 'TecMultiInterfering'),
('aec_single', 'AecSingleInterfering'),
]
def _find_tfrecord_records():
"""Finds the 4 benchmark records inside single_test_clean.tfrecord."""
by_clean = {}
for idx, raw in enumerate(tf.io.tf_record_iterator(TFRECORD_PATH)):
ex = tf.train.Example()
ex.ParseFromString(raw)
feat = ex.features.feature
clean_txt = feat['clean_transcript'].bytes_list.value[0].decode('utf-8')
int_txt = feat['interfering_transcript'].bytes_list.value[0].decode('utf-8')
utt_id = feat['utt_id'].bytes_list.value[0].decode('utf-8')
mixed_wav = np.array(
feat['mixed_waveform'].float_list.value, dtype=np.float32)
int_wav = np.array(
feat['interfering_waveform'].float_list.value, dtype=np.float32)
clean_wav = np.array(
feat['clean_waveform'].float_list.value, dtype=np.float32)
by_clean[clean_txt.strip()] = {
'tfrecord_index': idx,
'utt_id': utt_id,
'clean_txt': clean_txt,
'int_txt': int_txt,
'mixed_wav': mixed_wav,
'int_wav': int_wav,
'clean_wav': clean_wav,
}
return by_clean
def main():
os.makedirs(EXAMPLES_DIR, exist_ok=True)
by_clean = _find_tfrecord_records()
metadata = []
for item in TARGET_SAMPLES:
sid = item['id']
rec = by_clean.get(item['clean_ref'])
mixed_path = os.path.join(EXAMPLES_DIR, f'sample_{sid}_mixed.wav')
if rec is not None:
data_prep.write_wav_file(mixed_path, rec['mixed_wav'], sample_rate=24000)
mixed_wav, sr = data_prep.read_wav_file(mixed_path, target_sample_rate=24000)
if rec is not None:
int_wav = rec['int_wav']
utt_id = rec['utt_id']
tfr_idx = rec['tfrecord_index']
else:
lj_wav_path = (
'/usr/local/google/home/quanw/Data/tec_datasets/ljspeech/LJSpeech-1.1/wavs/'
f"{item['lj_id']}.wav"
)
int_wav, _ = data_prep.read_wav_file(lj_wav_path, target_sample_rate=24000)
if len(int_wav) < len(mixed_wav):
int_wav = np.pad(int_wav, (0, len(mixed_wav) - len(int_wav)))
else:
int_wav = int_wav[:len(mixed_wav)]
utt_id = item['utt_id']
tfr_idx = -1
int_path = os.path.join(EXAMPLES_DIR, f'sample_{sid}_interfering.wav')
data_prep.write_wav_file(int_path, int_wav, sample_rate=24000)
entry = {
'id': sid,
'utt_id': utt_id,
'tfrecord_index': tfr_idx,
'clean_ref': item['clean_ref'],
'interfering_text': item['interfering_text'],
'mixed_wav': f'examples/sample_{sid}_mixed.wav',
'interfering_wav': f'examples/sample_{sid}_interfering.wav',
'outputs': {},
}
for suffix, model_name in MODELS:
ckpt = os.path.join(CHECKPOINT_ROOT, model_name, 'best.ckpt')
out_rel = f'examples/sample_{sid}_{suffix}.wav'
out_abs = os.path.join(EXAMPLES_DIR, f'sample_{sid}_{suffix}.wav')
res = inference.run_inference(
mixed_waveform=mixed_wav,
model_name=model_name,
interfering_text=item['interfering_text'],
interfering_waveform=int_wav,
checkpoint_path=ckpt,
)
data_prep.write_wav_file(
out_abs, res['predicted_waveform'], sample_rate=sr)
entry['outputs'][suffix] = out_rel
if suffix == 'tec_single':
compat_abs = os.path.join(EXAMPLES_DIR, f'sample_{sid}_tec.wav')
data_prep.write_wav_file(
compat_abs, res['predicted_waveform'], sample_rate=sr)
print(f'Wrote {out_rel} ({len(res["predicted_waveform"])} samples)')
metadata.append(entry)
meta_path = os.path.join(EXAMPLES_DIR, 'metadata.json')
with open(meta_path, 'w', encoding='utf-8') as f:
json.dump(metadata, f, indent=2)
print(f'Saved metadata to {meta_path}')
if __name__ == '__main__':
main()
|