generate_examples.py

#!/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,
        'clean_ref': 'If I feel that way, I feel that way.',
        'interfering_text': 'Palm, as soon as he was seated in the cart,',
    },
    {
        'id': 2,
        'clean_ref': "And I've got nobody but you.",
        'interfering_text': 'Even now his sufferings were not at an end,',
    },
    {
        'id': 3,
        'clean_ref': '"Yes; it is perfection," she declared.',
        'interfering_text': (
            'Nearly all Americans are sensible and calm people.'
        ),
    },
    {
        'id': 4,
        'clean_ref': "If you come to me, I'll do as I see fit.",
        'interfering_text': (
            "A femoral pulse was also detected in the President's leg."
        ),
    },
]

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'])
    if rec is None:
      raise RuntimeError(f"Could not find record for {item['clean_ref']}")
    mixed_path = os.path.join(EXAMPLES_DIR, f'sample_{sid}_mixed.wav')
    if not os.path.exists(mixed_path):
      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)

    entry = {
        'id': sid,
        'utt_id': rec['utt_id'],
        'tfrecord_index': rec['tfrecord_index'],
        'clean_ref': item['clean_ref'],
        'interfering_text': item['interfering_text'],
        'mixed_wav': f'examples/sample_{sid}_mixed.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=rec['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()