initial commit

README.md

@@ -1,3 +1,132 @@
 ---
-license: apache-2.0
+language: "en"
+thumbnail:
+tags:
+- audio-classification
+- speechbrain
+- Emotion
+- Diarization
+- wavlm
+- pytorch
+license: "apache-2.0"
+datasets:
+- ZaionEmotionDataset
+- iemocap
+- ravdess
+- jl-corpus
+- esd
+- emov-db
+metrics:
+- EDER
 ---
+
+<iframe src="https://ghbtns.com/github-btn.html?user=speechbrain&repo=speechbrain&type=star&count=true&size=large&v=2" frameborder="0" scrolling="0" width="170" height="30" title="GitHub"></iframe>
+<br/><br/>
+
+# Emotion Diarization with WavLM Large on 5 popular emotional datasets.
+
+This repository provides all the necessary tools to perform speech emotion diarization with a fine-tuned wavlm (large) model using SpeechBrain.
+
+The model is trained on concatenated audios and tested on [ZaionEmotionDataset](https://zaion.ai/en/resources/zaion-lab-blog/zaion-emotion-dataset/). The metric is Emotion Diarization Error Rate (EDER). For more details please check the [paper link](https://arxiv.org/pdf/2306.12991.pdf).
+
+
+For a better experience, we encourage you to learn more about [SpeechBrain](https://speechbrain.github.io). The model performance on ZED (test set) is:
+
+| Release | EDER(%) | 
+|:-------------:|:--------------:|
+| 19-10-21 | 29.7 (Avg: 30.2) | 
+
+
+## Pipeline description
+
+This system is composed of an wavlm model. It is a combination of convolutional and residual blocks. The task aimes to predict the correct emotion composants and their boundaries within an utterance. For now, the model was trained with audios that contain only 1 non-neutral emotion event.
+
+The system is trained with recordings sampled at 16kHz (single channel).
+The code will automatically normalize your audio (i.e., resampling + mono channel selection) when calling *diarize_file* if needed.
+
+
+## Install SpeechBrain
+
+First of all, please install the **development** version of SpeechBrain with the following command:
+
+```
+pip install speechbrain
+```
+
+Please notice that we encourage you to read our tutorials and learn more about
+[SpeechBrain](https://speechbrain.github.io).
+
+### Perform Speech Emotion Diarization
+
+An external `py_module_file=custom.py` is used as an external Predictor class into this HF repos. We use `foreign_class` function from `speechbrain.pretrained.interfaces` that allow you to load you custom model. 
+
+```python
+from speechbrain.pretrained.interfaces import foreign_class
+classifier = foreign_class(
+    source="speechbrain/emotion-diarization-wavlm-large",
+    pymodule_file="custom_interface.py",
+    classname="Speech_Emotion_Diarization"
+)
+diary = classifier.diarize_file("speechbrain/emotion-diarization-wavlm-large/example.wav")
+print(diary)
+```
+The output will contain a dictionary of emotion composants and their boundaries.
+
+### Inference on GPU
+To perform inference on the GPU, add  `run_opts={"device":"cuda"}`  when calling the `from_hparams` method.
+
+### Training
+The model was trained with SpeechBrain (aa018540).
+To train it from scratch follows these steps:
+1. Clone SpeechBrain:
+```bash
+git clone https://github.com/speechbrain/speechbrain/
+```
+2. Install it:
+```
+cd speechbrain
+pip install -r requirements.txt
+pip install -e .
+```
+
+3. Run Training:
+```
+cd  recipes/ZaionEmotionDataset/emotion_diarization
+python train.py hparams/train.yaml --zed_folder /path/to/ZED --emovdb_folder /path/to/EmoV-DB --esd_folder /path/to/ESD --iemocap_folder /path/to/IEMOCAP --jlcorpus_folder /path/to/JL_corpus --ravdess_folder /path/to/RAVDESS
+```
+
+You can find our training results (models, logs, etc) [here](to be added).
+
+### Limitations
+The SpeechBrain team does not provide any warranty on the performance achieved by this model when used on other datasets.
+
+# **About Speech Emotion Diarization/Zaion Emotion Dataset**
+
+```bibtex
+@article{wang2023speech,
+  title={Speech Emotion Diarization: Which Emotion Appears When?},
+  author={Wang, Yingzhi and Ravanelli, Mirco and Nfissi, Alaa and Yacoubi, Alya},
+  journal={arXiv preprint arXiv:2306.12991},
+  year={2023}
+}
+```
+
+# **Citing SpeechBrain**
+Please, cite SpeechBrain if you use it for your research or business.
+
+```bibtex
+@misc{speechbrain,
+  title={{SpeechBrain}: A General-Purpose Speech Toolkit},
+  author={Mirco Ravanelli and Titouan Parcollet and Peter Plantinga and Aku Rouhe and Samuele Cornell and Loren Lugosch and Cem Subakan and Nauman Dawalatabad and Abdelwahab Heba and Jianyuan Zhong and Ju-Chieh Chou and Sung-Lin Yeh and Szu-Wei Fu and Chien-Feng Liao and Elena Rastorgueva and François Grondin and William Aris and Hwidong Na and Yan Gao and Renato De Mori and Yoshua Bengio},
+  year={2021},
+  eprint={2106.04624},
+  archivePrefix={arXiv},
+  primaryClass={eess.AS},
+  note={arXiv:2106.04624}
+}
+```
+
+# **About SpeechBrain**
+- Website: https://speechbrain.github.io/
+- Code: https://github.com/speechbrain/speechbrain/
+- HuggingFace: https://huggingface.co/speechbrain/

config.json

@@ -0,0 +1,3 @@
+{
+    "speechbrain_interface": "EncoderWav2vecClassifier"
+}

custom_interface.py

@@ -0,0 +1,215 @@
+import torch
+from speechbrain.pretrained import Pretrained
+
+class Speech_Emotion_Diarization(Pretrained):
+    """A ready-to-use SED interface (audio -> emotions and their durations)
+
+    Arguments
+    ---------
+    hparams
+        Hyperparameters (from HyperPyYAML)
+
+    Example
+    -------
+    >>> from speechbrain.pretrained import Speech_Emotion_Diarization
+    >>> tmpdir = getfixture("tmpdir")
+    >>> sed_model = Speech_Emotion_Diarization.from_hparams(source="speechbrain/emotion-diarization-wavlm-large", savedir=tmpdir,) # doctest: +SKIP
+    >>> sed_model.diarize_file("speechbrain/emotion-diarization-wavlm-large/example.wav") # doctest: +SKIP
+    """
+
+    MODULES_NEEDED = ["input_norm", "wav2vec", "output_mlp"]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def diarize_file(self, path):
+        """Get emotion diarization of a spoken utterance.
+
+        Arguments
+        ---------
+        path : str
+            Path to audio file which to diarize.
+
+        Returns
+        -------
+        dict
+            The emotions and their boundaries.
+        """
+        waveform = self.load_audio(path)
+        # Fake a batch:
+        batch = waveform.unsqueeze(0)
+        rel_length = torch.tensor([1.0])
+        frame_class = self.diarize_batch(
+            batch, rel_length, [path]
+        )
+        return frame_class
+
+    def encode_batch(self, wavs, wav_lens):
+        """Encodes audios into fine-grained emotional embeddings
+
+        Arguments
+        ---------
+        wavs : torch.tensor
+            Batch of waveforms [batch, time, channels].
+        wav_lens : torch.tensor
+            Lengths of the waveforms relative to the longest one in the
+            batch, tensor of shape [batch]. The longest one should have
+            relative length 1.0 and others len(waveform) / max_length.
+            Used for ignoring padding.
+
+        Returns
+        -------
+        torch.tensor
+            The encoded batch
+        """
+        if len(wavs.shape) == 1:
+            wavs = wavs.unsqueeze(0)
+
+        # Assign full length if wav_lens is not assigned
+        if wav_lens is None:
+            wav_lens = torch.ones(wavs.shape[0], device=self.device)
+
+        wavs, wav_lens = wavs.to(self.device), wav_lens.to(self.device)
+
+        wavs = self.mods.input_norm(wavs, wav_lens)
+        outputs = self.mods.wav2vec2(wavs)
+        return outputs
+
+
+    def diarize_batch(self, wavs, wav_lens, batch_id):
+        """Get emotion diarization of a batch of waveforms.
+
+        The waveforms should already be in the model's desired format.
+        You can call:
+        ``normalized = EncoderDecoderASR.normalizer(signal, sample_rate)``
+        to get a correctly converted signal in most cases.
+
+        Arguments
+        ---------
+        wavs : torch.tensor
+            Batch of waveforms [batch, time, channels].
+        wav_lens : torch.tensor
+            Lengths of the waveforms relative to the longest one in the
+            batch, tensor of shape [batch]. The longest one should have
+            relative length 1.0 and others len(waveform) / max_length.
+            Used for ignoring padding.
+
+        Returns
+        -------
+        torch.tensor
+            The frame-wise predictions
+        """
+        outputs = self.encode_batch(wavs, wav_lens)
+        averaged_out = self.hparams.avg_pool(outputs)
+        outputs = self.mods.output_mlp(averaged_out)
+        outputs = self.hparams.log_softmax(outputs)
+        score, index = torch.max(outputs, dim=-1)
+        preds = self.hparams.label_encoder.decode_torch(index)
+        results = self.preds_to_diarization(preds, batch_id)
+        return results
+
+    def preds_to_diarization(self, prediction, batch_id):
+        """Convert frame-wise predictions into a dictionary of
+        diarization results.
+
+        Returns
+        -------
+        dictionary
+            A dictionary with the start/end of each emotion
+        """
+        results = {}
+
+        for i in range(len(prediction)):
+            pred = prediction[i]
+            lol = []
+            for j in range(len(pred)):
+                start = round(self.hparams.stride * 0.02 * j, 2)
+                end = round(start + self.hparams.window_length * 0.02, 2)
+                lol.append([batch_id[i], start, end, pred[j]])
+
+            lol = merge_ssegs_same_emotion_adjacent(lol)
+            print(lol)
+            results[batch_id[i]] = [{"start": k[1], "end":k[2], "emotion": k[3]} for k in lol]
+            return results
+        
+        
+    def forward(self, wavs, wav_lens):
+        """Runs full transcription - note: no gradients through decoding"""
+        return self.transcribe_batch(wavs, wav_lens)
+
+
+def is_overlapped(end1, start2):
+    """Returns True if segments are overlapping.
+
+    Arguments
+    ---------
+    end1 : float
+        End time of the first segment.
+    start2 : float
+        Start time of the second segment.
+
+    Returns
+    -------
+    overlapped : bool
+        True of segments overlapped else False.
+
+    Example
+    -------
+    >>> from speechbrain.processing import diarization as diar
+    >>> diar.is_overlapped(5.5, 3.4)
+    True
+    >>> diar.is_overlapped(5.5, 6.4)
+    False
+    """
+
+    if start2 > end1:
+        return False
+    else:
+        return True
+
+
+def merge_ssegs_same_emotion_adjacent(lol):
+    """Merge adjacent sub-segs if they are the same emotion.
+    Arguments
+    ---------
+    lol : list of list
+        Each list contains [utt_id, sseg_start, sseg_end, emo_label].
+    Returns
+    -------
+    new_lol : list of list
+        new_lol contains adjacent segments merged from the same emotion ID.
+    Example
+    -------
+    >>> from speechbrain.utils.EDER import merge_ssegs_same_emotion_adjacent
+    >>> lol=[['u1', 0.0, 7.0, 'a'],
+    ... ['u1', 7.0, 9.0, 'a'],
+    ... ['u1', 9.0, 11.0, 'n'],
+    ... ['u1', 11.0, 13.0, 'n'],
+    ... ['u1', 13.0, 15.0, 'n'],
+    ... ['u1', 15.0, 16.0, 'a']]
+    >>> merge_ssegs_same_emotion_adjacent(lol)
+    [['u1', 0.0, 9.0, 'a'], ['u1', 9.0, 15.0, 'n'], ['u1', 15.0, 16.0, 'a']]
+    """
+    new_lol = []
+
+    # Start from the first sub-seg
+    sseg = lol[0]
+    flag = False
+    for i in range(1, len(lol)):
+        next_sseg = lol[i]
+        # IF sub-segments overlap AND has same emotion THEN merge
+        if is_overlapped(sseg[2], next_sseg[1]) and sseg[3] == next_sseg[3]:
+            sseg[2] = next_sseg[2]  # just update the end time
+            # This is important. For the last sseg, if it is the same emotion then merge
+            # Make sure we don't append the last segment once more. Hence, set FLAG=True
+            if i == len(lol) - 1:
+                flag = True
+                new_lol.append(sseg)
+        else:
+            new_lol.append(sseg)
+            sseg = next_sseg
+    # Add last segment only when it was skipped earlier.
+    if flag is False:
+        new_lol.append(lol[-1])
+
+    return new_lol

hyperparams.yaml

@@ -0,0 +1,68 @@
+# ############################################################################
+# Model: WavLM for Emotion Diarization
+# ############################################################################
+
+
+# Hparams NEEDED
+HPARAMS_NEEDED: ["window_length", "stride", "encoder_dim", "out_n_neurons", "avg_pool", "label_encoder", "softmax"]
+# Modules Needed
+MODULES_NEEDED: ["wav2vec2", "output_mlp"]
+
+# Feature parameters
+wav2vec2_hub: "microsoft/wavlm-large"
+
+# Pretrain folder (HuggingFace)
+pretrained_path: /home/ywang/zed_pr/sed_hf
+
+# parameters
+window_length: 1 # win_len = 0.02 * 1 = 0.02s
+stride: 1 # stride = 0.02 * 1 = 0.02s
+encoder_dim: 1024
+out_n_neurons: 4
+
+input_norm: !new:speechbrain.processing.features.InputNormalization
+    norm_type: sentence
+    std_norm: False
+
+wav2vec2: !new:speechbrain.lobes.models.huggingface_wav2vec.HuggingFaceWav2Vec2
+    source: !ref <wav2vec2_hub>
+    output_norm: True
+    freeze: False
+    freeze_feature_extractor: True
+    save_path: wav2vec2_checkpoint
+
+avg_pool: !new:speechbrain.nnet.pooling.Pooling1d
+    pool_type: "avg"
+    kernel_size: !ref <window_length>
+    stride: !ref <stride>
+    ceil_mode: True
+
+output_mlp: !new:speechbrain.nnet.linear.Linear
+    input_size: !ref <encoder_dim>
+    n_neurons: !ref <out_n_neurons>
+    bias: False
+
+model: !new:torch.nn.ModuleList
+    - [!ref <output_mlp>]
+
+modules:
+    input_norm: !ref <input_norm>
+    wav2vec2: !ref <wav2vec2>
+    output_mlp: !ref <output_mlp>
+
+log_softmax: !new:speechbrain.nnet.activations.Softmax
+    apply_log: True
+
+label_encoder: !new:speechbrain.dataio.encoder.CategoricalEncoder
+
+pretrainer: !new:speechbrain.utils.parameter_transfer.Pretrainer
+    loadables:
+        input_norm: !ref <input_norm>
+        wav2vec2: !ref <wav2vec2>
+        model: !ref <model>
+        label_encoder: !ref <label_encoder>
+    paths:
+        input_norm: !ref <pretrained_path>/input_norm.ckpt
+        wav2vec2: !ref <pretrained_path>/wav2vec2.ckpt
+        model: !ref <pretrained_path>/model.ckpt
+        label_encoder: !ref <pretrained_path>/label_encoder.txt

input_norm.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eddbd59b97a6456c5a81880065b785f731ca3b959abfa2c965658a591e53d31f
+size 1075

label_encoder.txt

@@ -0,0 +1,6 @@
+'a' => 0
+'n' => 1
+'h' => 2
+'s' => 3
+================
+'starting_index' => 0

model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:23c5832103c64cb628e8e56ce5fc7061be323e435a294d34060172c10015208d
+size 17189

wav2vec2.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:60d8746058853c9ad8976c5630d8584959b74428f6bbe7458fe3d8bdf15d54b3
+size 1262005979