Add HuBERT-fbeeper demo

fbeeper_hubert.py

@@ -0,0 +1,157 @@
+"""Beep profanity words in audio using one of the Hubert-compatible ASR models.
+"""
+
+import argparse
+import re
+import logging
+
+import soundfile
+import transformers
+import torch
+import numpy as np
+
+
+log = logging.getLogger(__name__)
+
+
+class HubertBeeper:
+
+    PROFANITY = ["fuck", "shit", "piss"]
+
+    def __init__(self, model_name="facebook/hubert-large-ls960-ft"):
+        log.debug("Loading model: %s", model_name)
+        self.model_name = model_name
+        self.model = transformers.AutoModelForCTC.from_pretrained(model_name)
+        self.model.eval()
+        self.feature_extractor = transformers.AutoFeatureExtractor.from_pretrained(model_name)
+        self.tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
+        self.processor = transformers.Wav2Vec2Processor(
+            feature_extractor=self.feature_extractor, tokenizer=self.tokenizer)
+
+    def asr(self, waveform, sample_rate):
+        features = self.processor([waveform], sampling_rate=sample_rate)
+        features = torch.tensor(features.input_values)
+        output = self.model(features)
+        return output
+
+    def f_beep(self, sound_file_path: str) -> np.array:
+        wav, sample_rate = soundfile.read(sound_file_path)
+        text, result_wav = self.f_beep_waveform(wav, sample_rate)
+        return result_wav
+
+    def f_beep_waveform(self, wav: np.array, sample_rate: int) -> np.array:
+        model_output = self.asr(wav, sample_rate)
+        text, spans = find_words_in_audio(model_output, self.processor, self.PROFANITY)
+        number_of_frames = model_output.logits.shape[1]
+        frame_size = len(wav) / number_of_frames
+
+        # Mask offsensive parts of the audio
+        for frame_begin, frame_end in spans:
+            begin = round(frame_begin * frame_size)
+            end = round(frame_end * frame_size)
+            self.generate_beep(wav, begin, end)
+
+        return text, wav
+
+    def generate_beep(self, wav, begin, end):
+        """Generate a beep over the selected region in audio.
+
+        Modifies waveform in place.
+        """
+
+        # Silence sounds better than beeps
+        for i in range(begin, end):
+            wav[i] = 0
+
+def find_words_in_audio(model_output, processor, words):
+    """Return all frame spans that matches any of the `words`.
+    """
+
+    result_spans = []
+
+    token_ids = model_output.logits.argmax(dim=-1)[0]
+    vocab = processor.tokenizer.get_vocab()
+    text, offsets = decode_output_with_offsets(token_ids, vocab)
+    text = text.lower()
+    log.debug("ASR text: %s", text)
+
+    for word in words:
+        result_spans += find_spans(text, offsets, word)
+
+    log.debug("Spans: %s", result_spans)
+
+    return text, result_spans
+
+
+def find_spans(text, offsets, word):
+    """Return all frame indexes that correspond to the given `word`.
+    """
+    spans = []
+    pattern = r"\b" + re.escape(word) + r"\b"
+    for match in re.finditer(pattern, text):
+        a = match.start() 
+        b = match.end() + 1
+        start_frame = offsets[a]
+        end_frame = offsets[b] if b < len(offsets) else -1
+        spans.append((start_frame, end_frame))
+    return spans
+
+
+def decode_output_with_offsets(decoded_token_ids, vocab):
+    """Given list of decoded tokens, return text and
+    time offsets that correspond to each character in the text.
+    
+    Args:
+        decoded_token_ids (List[int]): list of token ids.
+            The length of the list should be equal to the number
+            of audio frames.
+        vocab (Dict[str, int]): model's vocabulary.
+        
+    Returns:
+        Tuple[str, List[int]], where
+            `str` is a decoded text,
+            `List[int]` is a starting frame indexes for
+            every character in text.
+    """
+    token_by_index = {v: k for k, v in vocab.items()}
+    prev_token = None
+    result_string = []
+    result_offset = []
+    for i, token_id in enumerate(decoded_token_ids):
+        token_id = token_id.item()
+        if token_id == 0:
+            continue
+            
+        token = token_by_index[token_id]
+        if token == prev_token:
+            continue
+        
+        result_string.append(token)
+        result_offset.append(i)
+        prev_token = token
+        
+    result_string = "".join(result_string).replace("|", " ")
+    assert len(result_string) == len(result_offset)
+    return result_string, result_offset
+
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("input")
+    parser.add_argument("-o", "--output")
+    parser.add_argument("-v", "--verbose", action="store_true")
+    parser.add_argument("--model", default="facebook/hubert-large-ls960-ft")
+    args = parser.parse_args()
+    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)
+
+    beeper = HubertBeeper(args.model)
+    result = beeper.f_beep(args.input)
+
+    output = args.output or "result.wav"
+    soundfile.write(output, result, 16000)
+    print(f"Saved to {output}")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,7 @@
+transformers==4.12.3
+pydub
+soundfile
+librosa
+unidecode
+inflect
+torchaudio