DOCSEARCHbyAlgolia
Submit Docs

AV-HuBERT

repository·main·Indexed 21 days ago

https://github.com/facebookresearch/av_hubert

A self-supervised representation learning framework for audio-visual speech, used for tasks such as lip reading, Automatic Speech Recognition (ASR), and audio-visual speech recognition. The repository provides tools for data preprocessing of LRS3 and VoxCeleb2 datasets, MFCC and transformer layer feature extraction, K-means clustering for label generation, and model pre-training and fine-tuning using Fairseq.

Tokens
4.6K
Snippets
15
Records
17
Agent score
75%

What's inside AV-HuBERT

  1. Prepare TSV files for label generation

    main

    The scripts expect *.tsv files containing audio metadata. Each file must follow this format:

    • Line 1: <root-dir>
    • Subsequent lines: <id> <video-path> <audio-path> <video-number-frames> <audio-number-frames> (tab-separated).

    Example structure:

    <root-dir>
<id-1> <video-path-1> <audio-path-1> <video-number-frames-1> <audio-number-frames-1>
<id-2> <video-path-2> <audio-path-2> <video-number-frames-2> <audio-number-frames-2>
...```

  2. Merge K-means label shards and create cluster dictionary

    main

    After generating labels for all shards, you must merge them and create a dictionary of cluster indexes.

    1. Merge shards:

    for rank in $(seq 0 $((nshard - 1))); do
  cat $lab_dir/${split}_${rank}_${nshard}.km
done > $lab_dir/${split}.km

    2. Create cluster index dictionary: (Note: This assumes n_cluster is the number of clusters)

    for i in $(seq 1 $((n_cluster-1)));do 
    echo $i 10000
done > $lab_dir/dict.{mfcc,km}.txt

  4. Install AV-HuBERT

    main

    To set up the AV-HuBERT environment, follow these steps to create a conda environment, clone the repository with submodules, and install dependencies including Fairseq.

    1. Create and activate a conda environment with Python 3.8.
    2. Clone the repository and initialize submodules.
    3. Install requirements and install Fairseq in editable mode.
    # Create and activate conda environment
conda create -n avhubert python=3.8 -y
conda activate avhubert

# Clone repository and submodules
git clone https://github.com/facebookresearch/av_hubert.git
cd avhubert
git submodule init
git submodule update

# Install dependencies
pip install -r requirements.txt
cd fairseq
pip install --editable ./

  6. Prepare audio noise for environment simulation

    main

    To simulate noisy environments, you can prepare noise data using the MUSAN corpus and LRS3.

    1. MUSAN Preparation

    Download MUSAN (folders: music, speech, noise) and run musan_prepare.py. This splits audio into 10s clips, generates babble noise from speech, and counts frames. Requires Slurm/submitit.

    2. LRS3 Noise Preparation

    • Generate LRS3 babble and speech noise: python lrs3_noise.py --lrs3 ${lrs3}. Results are in ${lrs3}/noise/{babble,speech}.
    • Generate babble noise specifically from the LRS3 training set: python mix_babble.py --lrs3 ${lrs3}.

    3. Noise Manifest Generation

    Run noise_manifest.py --lrs3 ${lrs3} --musan ${musan} to create .tsv manifests for MUSAN babble, music, and noise, as well as a combined manifest in ${musan}/tsv/all.

    # 1. MUSAN data preparation
python musan_prepare.py --musan ${musan} --nshard ${nshard} --slurm_partition ${slurm_partition}

# 2. LRS3 audio noise preparation
python lrs3_noise.py --lrs3 ${lrs3}
python mix_babble.py --lrs3 ${lrs3}

# 3. Set up noise directory
python noise_manifest.py --lrs3 ${lrs3} --musan ${musan}

  7. Preprocess LRS3 dataset

    main

    Preprocessing the LRS3 dataset involves four main stages. Assume the LRS3 data directory is ${lrs3} and contains pretrain, trainval, and test folders.

    1. Data Preparation

    Run lrs3_prepare.py to generate file lists and text labels. The --step flag determines the operation:

    • Step 1: Splits long utterances in pretraining into shorter ones and generates time boundaries/labels.
    • Step 2: Trims videos and audios based on new boundaries (uses --rank and --nshard for sharding).
    • Step 3: Extracts audio for trainval and test splits (uses --rank and --nshard).
    • Step 4: Generates file IDs and text transcriptions.

    2. Facial Landmark Detection and Mouth ROI Cropping

    This stage uses dlib to detect landmarks and align_mouth.py to crop the mouth regions. This generates mouth ROIs in ${lrs3}/video.

    3. Frame Counting

    Use count_frames.py to count audio/video frames. For Slurm users, count_frames_slurm.py can automate sharding and merging.

    4. Manifest Generation

    Use lrs3_manifest.py to set up the final data directory (e.g., for trainval-only or pretrain+trainval sets). This creates a tokenizer and target directories containing .tsv (manifests) and .wrd (text labels) files.

    # 1. Data preparation
python lrs3_prepare.py --lrs3 ${lrs3} --ffmpeg /path/to/ffmpeg --rank ${rank} --nshard ${nshard} --step ${step}

# 2. Detect facial landmark and crop mouth ROIs
python detect_landmark.py --root ${lrs3} --landmark ${lrs3}/landmark --manifest ${lrs3}/file.list \
 --cnn_detector /path/to/dlib_cnn_detector --face_detector /path/to/dlib_landmark_predictor --ffmpeg /path/to/ffmpeg \
 --rank ${rank} --nshard ${nshard}

python align_mouth.py --video-direc ${lrs3} --landmark ${landmark_dir} --filename-path ${lrs3}/file.list \
 --save-direc ${lrs3}/video --mean-face /path/to/mean_face --ffmpeg /path/to/ffmpeg \
 --rank ${rank} --nshard ${nshard}

# 3. Count number of frames per clip
python count_frames.py --root ${lrs3} --manifest ${lrs3}/file.list --nshard ${nshard} --rank ${rank}

# 4. Set up data directory
python lrs3_manifest.py --lrs3 ${lrs3} --manifest ${lrs3}/file.list \
 --valid-ids /path/to/valid --vocab-size ${vocab_size}

  8. Preprocess VoxCeleb2 dataset

    main

    Preprocessing VoxCeleb2 follows a similar pattern to LRS3. Assume the data directory is ${vox} and contains dev and test folders.

    1. Data Preparation

    Run vox_prepare.py. --step 1 generates file IDs, and --step 2 extracts audio waveforms (uses --rank and --nshard).

    2. Facial Landmark Detection and Mouth ROI Cropping

    Use detect_landmark.py and align_mouth.py to generate mouth ROIs in ${vox}/video.

    3. Frame Counting

    Use count_frames.py to count frames. For Slurm, use count_frames_slurm.py.

    4. Manifest Generation

    Use vox_manifest.py to set up the data directory for the full dataset or an English-only subset (using --en-ids).

    # 1. Data preparation
python vox_prepare.py --root ${vox} --ffmpeg /path/to/ffmpeg --rank ${rank} --nshard ${nshard} --step ${step}

# 2. Detect facial landmark and crop mouth ROIs
python detect_landmark.py --root ${vox} --landmark ${vox}/landmark --manifest ${vox}/file.list \
 --cnn_detector /path/to/dlib_cnn_detector --face_detector /path/to/dlib_landmark_predictor --ffmpeg /path/to/ffmpeg \
 --rank ${rank} --nshard ${nshard}

python align_mouth.py --video-direc ${vox} --landmark ${landmark_dir} --filename-path ${vox}/file.list \
 --save-direc ${vox}/video --mean-face /path/to/mean_face --ffmpeg /path/to/ffmpeg \
 --rank ${rank} --nshard ${nshard}

# 3. Count number of frames per clip
python count_frames.py --root ${vox} --manifest ${vox}/file.list --nshard ${nshard} --rank ${rank}

# 4. Set up data directory
python vox_manifest.py --vox ${vox} --manifest ${vox}/file.list --en-ids /path/to/en

  9. Fit a K-means model

    main

    Use learn_kmeans.py to fit a K-means model on a subset of the extracted features.

    Arguments:

    • ${feat_dir}: Directory containing extracted features.
    • ${split}: The data split.
    • ${nshard}: Total number of shards.
    • ${km_path}: Path where the K-means model will be saved.
    • ${n_cluster}: Number of clusters to fit.
    • --percent: The fraction of data to use for fitting (e.g., 0.1 for 10%). Set to -1 to use all available data.

    Additional K-means options can be accessed via the -h flag.

    python learn_kmeans.py ${feat_dir} ${split} ${nshard} ${km_path} ${n_cluster} --percent 0.1

  10. Extract MFCC features

    main

    To extract 39-D mfcc+delta+ddelta features for the 1st iteration AV-HuBERT training, use dump_mfcc_feature.py. The process shards the TSV file into ${nshard} parts, and you run the command for each ${rank} in the range [0, nshard-1].

    Features are saved at ${feat_dir}/${split}_${rank}_${nshard}.{npy,len}.

    python dump_mfcc_feature.py ${tsv_dir} ${split} ${nshard} ${rank} ${feat_dir}

  11. Extract AV-HuBERT features from a trained model

    main

    To extract features from a specific transformer layer of a trained AV-HuBERT model, use dump_hubert_feature.py.

    Arguments:

    • ${tsv_dir}: Directory containing TSV files.
    • ${split}: The data split (e.g., train, valid).
    • ${ckpt_path}: Path to the trained AV-HuBERT checkpoint.
    • ${layer}: The index of the transformer layer to extract from.
    • ${nshard}: Total number of shards.
    • ${rank}: The current shard index [0, nshard-1].
    • ${feat_dir}: Directory to save features.
    • --user_dir: Path to the user directory (use pwd/../ to point to the project root).
    • --max_chunk: Use this flag to decrease chunk size if you encounter out-of-memory errors.

    Features are saved at ${feat_dir}/${split}_${rank}_${nshard}.{npy,len}.

    python dump_hubert_feature.py ${tsv_dir} ${split} ${ckpt_path} ${layer} ${nshard} ${rank} ${feat_dir} --user_dir `pwd`/../

  12. Apply K-means model to generate labels

    main

    Use dump_km_label.py to apply a trained K-means model ${km_path} to the features of a specific ${split}.

    Arguments:

    • ${feat_dir}: Directory containing extracted features.
    • ${split}: The data split.
    • ${km_path}: Path to the trained K-means model.
    • ${nshard}: Total number of shards.
    • ${rank}: The current shard index [0, nshard-1].
    • ${lab_dir}: Directory to save the resulting labels.

    Labels for the ${rank}-th shard are dumped to ${lab_dir}/${split}_${rank}_${nshard}.km.

    python dump_km_label.py ${feat_dir} ${split} ${km_path} ${nshard} ${rank} ${lab_dir}