This repository provides the data and source code for Text-to-Audio Grounding (TAG) task.
We provide a quick inference usage with Hugging Face:
import torch
import torchaudio
from transformers import AutoModel
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = AutoModel.from_pretrained(
"wsntxxn/cnn8rnn-w2vmean-audiocaps-grounding",
trust_remote_code=True
).to(device)
wav1, sr1 = torchaudio.load("/path/to/file1.wav")
wav1 = torchaudio.functional.resample(wav1, sr1, model.config.sample_rate)
wav1 = wav1.mean(0) if wav1.size(0) > 1 else wav1[0]
wav2, sr2 = torchaudio.load("/path/to/file2.wav")
wav2 = torchaudio.functional.resample(wav2, sr2, model.config.sample_rate)
wav2 = wav2.mean(0) if wav2.size(0) > 1 else wav2[0]
wav_batch = torch.nn.utils.rnn.pad_sequence([wav1, wav2], batch_first=True)
text = ["a man speaks", "a dog is barking"]
with torch.no_grad():
output = model(
audio=wav_batch,
audio_length=[wav1.size(0), wav2.size(0)],
text=text
)
# output: (2, n_seconds // 4)
output
gives the estimated probability of the input text prompt.
The AudioGrounding dataset is an augmented audio captioning dataset. It is based on AudioCaps, which is established using part of a audio event dataset, AudioSet. Therefore, audio files can be downloaded from AudioSet.
The updated AudioGrounding v2 is available in .
Changes in version 2:
- Train/val/test sets are re-split and refined.
- Data are re-formatted and audio files are renamed.
The current label format: a list of audio_item
, each containing
- audiocap_id:
id
in the AudioCaps label file - audio_id: audio filename, in the form of
Y[youtube_id].wav
- tokens: caption tokenized by NLTK
- phrases: a list of
phrase_item
- phrase: tokens of the current query phrase
- start_index: index of the first token of
phrase
, starting from 0 - end_index: index of the last token of
phrase
- segments: a list of
[onset, offset]
timestamp annotations
We provide a baseline approach for TAG in this repository. To run the baseline:
- checkout the code and install the required python packages:
git clone https://github.com/wsntxxn/TextToAudioGrounding
pip install -r requirements.txt
- download audio clips and labels from Zenodo.
- pack waveforms, assume the audio files are in
$AUDIO
:
mkdir data/audiogrounding
for split in train val test; do
python utils/data/prepare_wav_csv.py $AUDIO/$split data/audiogrounding/$split/wav.csv
python utils/data/pack_waveform.py data/audiogrounding/$split/wav.csv \
-o data/audiogrounding/$split/waveform.h5 \
--sample_rate 32000
done
python utils/data/prepare_duration.py data/audiogrounding/test/wav.csv data/audiogrounding/test/duration.csv
- prepare vocabulary file:
python utils/build_vocab.py data/audiogrounding/train/label.json data/audiogrounding/train/vocab.pkl
- run the training and evaluation:
python python_scripts/training/run_strong.py train_evaluate \
--train_config $TRAIN_CFG \
--eval_config $EVAL_CFG
Or alternatively,
python python_scripts/training/run_strong.py train \
--config $TRAIN_CFG
python python_scripts/training/run_strong.py evaluate \
--experiment_path $EXP_PATH \
--eval_config $EVAL_CFG
$TRAIN_CFG
and $EVAL_CFG
are yaml-formatted configuration files.
$EXP_PATH
is the checkpoint directory set in $TRAIN_CFG
.
Example configuration files are provided here.
We provide the best-performing WSTAG model, downloaded here. Unzip it into $MODEL_DIR
:
unzip audiocaps_cnn8rnn_w2vmean_dp_ls_clustering_selfsup.zip -d $MODEL_DIR
Remember to modify the training data vocabulary path in $MODEL_DIR/config.yaml
(data.train.collate_fn.tokenizer.args.vocabulary) to $MODEL_DIR/vocab.pkl
.
To ensure that the vocabulary file used for training is loaded for inference, the inference script uses the vocabulary path specified in $MODEL_DIR/config.yaml
.
Inference:
python python_scripts/inference/inference.py inference_multi_text_model \
--experiment_path $MODEL_DIR \
--audio $AUDIO \
--phrase $PHRASE \
--output ./prob.png
For all settings, training is done in the same way as in the baseline:
python $TRAIN_SCRIPT train_evaluate \
--train_config $TRAIN_CFG \
--eval_config $EVAL_CFG
The training scripts and configurations vary for different settings. We provide the training script and example configuration file in each setting.
WSTAG uses audio captioning data for training.
The format of training data is the same as AudioGrounding, with the only difference that there is no segments
in phrase_item
.
You can convert the original captioning data into this format by yourself.
The phrase parsing rules are provided here.
The waveform packing and vocabulary preparation process is also the same as in the baseline.
TRAIN_SCRIPT
: run_weak_sentence.pyTRAIN_CFG
: cnn8rnn_w2vmean_dp_amean_tmean.yamlEVAL_CFG
: eval.yaml
For all phrase-level settings, the example EVAL_CFG
is eval.yaml.
For all phrase-level settings except "X + self-supervision", TRAIN_SCRIPT
is run_weak_phrase.py.
TRAIN_CFG
: cnn8rnn_w2vmean_random.yaml
TRAIN_CFG
: cnn8rnn_w2vmean_similarity.yaml
Similarity-based sampling requires pre-computed phrase embeddings.
We use the contrastive audio-text model trained on AudioCaps to extract phrase embeddings.
Download the model from here, unzip it into $CLAP_DIR
, then extract embeddings:
unzip audiocaps_cnn14_bertm.zip -d $CLAP_DIR
python utils/data/create_text_embedding/prepare_phrase_clap.py phrase \
--experiment_path $CLAP_DIR \
--phrase_input $DATA \
--output $OUTPUT \
--with_proj True
Then modify the data.train.dataset.args.phrase_embed item in the training configuration file to $OUTPUT
accordingly.
TRAIN_CFG
: cnn8rnn_w2vmean_clustering.yaml
Clustering-based sampling requires clustering models. We train clustering models based on the pre-computed phrase embeddings.
python python_scripts/clustering/kmeans_emb.py \
--embedding $PHRASE_EMB \
--n_cluster $N_CLUSTER \
--output $OUTPUT
$PHRASE_EMB
is the phrase embedding file, i.e., $OUTPUT
of the previous step.
Remember to modify the data.train.dataset.args.cluster_map to the corresponding mapping file (not exactly $OUTPUT
).
TRAIN_SCRIPT
: run_weak_phrase_self_supervision.pyTRAIN_CFG
: cnn8rnn_w2vmean_clustering_selfsup.yaml
The teacher.pretrained should be set to the checkpoint path of the pretrained WSTAG model.
If you find the data useful, please cite our paper:
@inproceedings{xu2021text,
title={Text-to-audio grounding: Building correspondence between captions and sound events},
author={Xu, Xuenan and Dinkel, Heinrich and Wu, Mengyue and Yu, Kai},
booktitle={IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
pages={606--610},
year={2021},
organization={IEEE}
}
If you find the pre-trained WSTAG model useful, please cite this paper:
@article{xu2024towards,
title={Towards Weakly Supervised Text-to-Audio Grounding},
author={Xu, Xuenan and Ma, Ziyang and Wu, Mengyue and Yu, Kai},
journal={arXiv preprint arXiv:2401.02584},
year={2024}
}