src/audio/processing.py

# Standard library imports
import os
import re
import json
from io import TextIOWrapper
from typing import Annotated, Optional, Tuple, List, Dict

# Related third party imports
import torch
import faster_whisper
from pydub import AudioSegment
from deepmultilingualpunctuation import PunctuationModel

# Local imports
from src.audio.utils import TokenizerUtils


class AudioProcessor:
    """
    A class to handle various audio processing tasks, such as conversion,
    trimming, merging, and audio transformations.

    Parameters
    ----------
    audio_path : str
        Path to the audio file to process.
    temp_dir : str, optional
        Directory for storing temporary files. Defaults to ".temp".

    Attributes
    ----------
    audio_path : str
        Path to the input audio file.
    temp_dir : str
        Path to the temporary directory for processed files.
    mono_audio_path : Optional[str]
        Path to the mono audio file after conversion.

    Methods
    -------
    convert_to_mono()
        Converts the audio file to mono.
    get_duration()
        Gets the duration of the audio file in seconds.
    change_format(new_format)
        Converts the audio file to a new format.
    trim_audio(start_time, end_time)
        Trims the audio file to the specified time range.
    adjust_volume(change_in_db)
        Adjusts the volume of the audio file.
    get_channels()
        Gets the number of audio channels.
    fade_in_out(fade_in_duration, fade_out_duration)
        Applies fade-in and fade-out effects to the audio.
    merge_audio(other_audio_path)
        Merges the current audio with another audio file.
    split_audio(chunk_duration)
        Splits the audio file into chunks of a specified duration.
    create_manifest(manifest_path)
        Creates a manifest file containing metadata about the audio.
    """

    def __init__(
            self,
            audio_path: Annotated[str, "Path to the audio file"],
            temp_dir: Annotated[str, "Directory for temporary processed files"] = ".temp"
    ) -> None:
        if not isinstance(audio_path, str):
            raise TypeError("Expected 'audio_path' to be a string.")
        if not isinstance(temp_dir, str):
            raise TypeError("Expected 'temp_dir' to be a string.")

        self.audio_path = audio_path
        self.temp_dir = temp_dir
        self.mono_audio_path = None
        os.makedirs(temp_dir, exist_ok=True)

    def convert_to_mono(self) -> Annotated[str, "Path to the mono audio file"]:
        """
        Convert the audio file to mono.

        Returns
        -------
        str
            Path to the mono audio file.

        Examples
        --------
        >>> processor = AudioProcessor("example.wav")
        >>> mono_path = processor.convert_to_mono()
        >>> isinstance(mono_path, str)
        True
        """
        sound = AudioSegment.from_file(self.audio_path)
        mono_sound = sound.set_channels(1)
        self.mono_audio_path = os.path.join(self.temp_dir, "mono_file.wav")
        mono_sound.export(self.mono_audio_path, format="wav")
        return self.mono_audio_path

    def get_duration(self) -> Annotated[float, "Audio duration in seconds"]:
        """
        Get the duration of the audio file.

        Returns
        -------
        float
            Duration of the audio in seconds.

        Examples
        --------
        >>> processor = AudioProcessor("example.wav")
        >>> duration = processor.get_duration()
        >>> isinstance(duration, float)
        True
        """
        sound = AudioSegment.from_file(self.audio_path)
        return len(sound) / 1000.0

    def change_format(
            self, new_format: Annotated[str, "New audio format"]
    ) -> Annotated[str, "Path to converted audio file"]:
        """
        Convert the audio file to a new format.

        Parameters
        ----------
        new_format : str
            Desired format for the output audio file.

        Returns
        -------
        str
            Path to the converted audio file.

        Examples
        --------
        >>> processor = AudioProcessor("example.wav")
        >>> converted_path = processor.change_format("mp3")
        >>> isinstance(converted_path, str)
        True
        """
        if not isinstance(new_format, str):
            raise TypeError("Expected 'new_format' to be a string.")

        sound = AudioSegment.from_file(self.audio_path)
        output_path = os.path.join(self.temp_dir, f"converted_file.{new_format}")
        sound.export(output_path, format=new_format)
        return output_path

    def trim_audio(
            self, start_time: Annotated[float, "Start time in seconds"],
            end_time: Annotated[float, "End time in seconds"]
    ) -> Annotated[str, "Path to trimmed audio file"]:
        """
        Trim the audio file to the specified duration.

        Parameters
        ----------
        start_time : float
            Start time in seconds.
        end_time : float
            End time in seconds.

        Returns
        -------
        str
            Path to the trimmed audio file.

        Examples
        --------
        >>> processor = AudioProcessor("example.wav")
        >>> trimmed_path = processor.trim_audio(0.0, 10.0)
        >>> isinstance(trimmed_path, str)
        True
        """
        if not isinstance(start_time, (int, float)):
            raise TypeError("Expected 'start_time' to be a float or int.")
        if not isinstance(end_time, (int, float)):
            raise TypeError("Expected 'end_time' to be a float or int.")

        sound = AudioSegment.from_file(self.audio_path)
        trimmed_audio = sound[start_time * 1000:end_time * 1000]
        trimmed_audio_path = os.path.join(self.temp_dir, "trimmed_file.wav")
        trimmed_audio.export(trimmed_audio_path, format="wav")
        return trimmed_audio_path

    def adjust_volume(
            self, change_in_db: Annotated[float, "Volume change in dB"]
    ) -> Annotated[str, "Path to volume-adjusted audio file"]:
        """
        Adjust the volume of the audio file.

        Parameters
        ----------
        change_in_db : float
            Volume change in decibels.

        Returns
        -------
        str
            Path to the volume-adjusted audio file.

        Examples
        --------
        >>> processor = AudioProcessor("example.wav")
        >>> adjusted_path = processor.adjust_volume(5.0)
        >>> isinstance(adjusted_path, str)
        True
        """
        if not isinstance(change_in_db, (int, float)):
            raise TypeError("Expected 'change_in_db' to be a float or int.")

        sound = AudioSegment.from_file(self.audio_path)
        adjusted_audio = sound + change_in_db
        adjusted_audio_path = os.path.join(self.temp_dir, "adjusted_volume.wav")
        adjusted_audio.export(adjusted_audio_path, format="wav")
        return adjusted_audio_path

    def get_channels(self) -> Annotated[int, "Number of channels"]:
        """
        Get the number of audio channels.

        Returns
        -------
        int
            Number of audio channels.

        Examples
        --------
        >>> processor = AudioProcessor("example.wav")
        >>> channels = processor.get_channels()
        >>> isinstance(channels, int)
        True
        """
        sound = AudioSegment.from_file(self.audio_path)
        return sound.channels

    def fade_in_out(
            self, fade_in_duration: Annotated[float, "Fade-in duration in seconds"],
            fade_out_duration: Annotated[float, "Fade-out duration in seconds"]
    ) -> Annotated[str, "Path to faded audio file"]:
        """
        Apply fade-in and fade-out effects to the audio file.

        Parameters
        ----------
        fade_in_duration : float
            Duration of the fade-in effect in seconds.
        fade_out_duration : float
            Duration of the fade-out effect in seconds.

        Returns
        -------
        str
            Path to the faded audio file.

        Examples
        --------
        >>> processor = AudioProcessor("example.wav")
        >>> faded_path = processor.fade_in_out(1.0, 2.0)
        >>> isinstance(faded_path, str)
        True
        """
        if not isinstance(fade_in_duration, (int, float)):
            raise TypeError("Expected 'fade_in_duration' to be a float or int.")
        if not isinstance(fade_out_duration, (int, float)):
            raise TypeError("Expected 'fade_out_duration' to be a float or int.")

        sound = AudioSegment.from_file(self.audio_path)
        faded_audio = sound.fade_in(fade_in_duration * 1000).fade_out(fade_out_duration * 1000)
        faded_audio_path = os.path.join(self.temp_dir, "faded_audio.wav")
        faded_audio.export(faded_audio_path, format="wav")
        return faded_audio_path

    def merge_audio(
            self, other_audio_path: Annotated[str, "Path to other audio file"]
    ) -> Annotated[str, "Path to merged audio file"]:
        """
        Merge the current audio file with another audio file.

        Parameters
        ----------
        other_audio_path : str
            Path to the other audio file.

        Returns
        -------
        str
            Path to the merged audio file.

        Examples
        --------
        >>> processor = AudioProcessor("example.wav")
        >>> merged_path = processor.merge_audio("other_example.wav")
        >>> isinstance(merged_path, str)
        True
        """
        if not isinstance(other_audio_path, str):
            raise TypeError("Expected 'other_audio_path' to be a string.")

        sound1 = AudioSegment.from_file(self.audio_path)
        sound2 = AudioSegment.from_file(other_audio_path)
        merged_audio = sound1 + sound2
        merged_audio_path = os.path.join(self.temp_dir, "merged_audio.wav")
        merged_audio.export(merged_audio_path, format="wav")
        return merged_audio_path

    def split_audio(
            self, chunk_duration: Annotated[float, "Chunk duration in seconds"]
    ) -> Annotated[List[str], "Paths to audio chunks"]:
        """
        Split the audio file into chunks of the specified duration.

        Parameters
        ----------
        chunk_duration : float
            Duration of each chunk in seconds.

        Returns
        -------
        List[str]
            Paths to the generated audio chunks.

        Examples
        --------
        >>> processor = AudioProcessor("example.wav")
        >>> chunks = processor.split_audio(10.0)
        >>> isinstance(chunks, list)
        True
        """
        if not isinstance(chunk_duration, (int, float)):
            raise TypeError("Expected 'chunk_duration' to be a float or int.")

        sound = AudioSegment.from_file(self.audio_path)
        chunk_paths = []

        for i in range(0, len(sound), int(chunk_duration * 1000)):
            chunk = sound[i:i + int(chunk_duration * 1000)]
            chunk_path = os.path.join(self.temp_dir, f"chunk_{i // 1000}.wav")
            chunk.export(chunk_path, format="wav")
            chunk_paths.append(chunk_path)

        return chunk_paths

    def create_manifest(
            self,
            manifest_path: Annotated[str, "Manifest file path"]
    ) -> None:
        """
        Create a manifest file containing metadata about the audio file.

        Parameters
        ----------
        manifest_path : str
            Path to the manifest file.

        Examples
        --------
        >>> processor = AudioProcessor("example.wav")
        >>> processor.create_manifest("manifest.json")
        """
        duration = self.get_duration()
        manifest_entry = {
            "audio_filepath": self.audio_path,
            "offset": 0,
            "duration": duration,
            "label": "infer",
            "text": "-",
            "rttm_filepath": None,
            "uem_filepath": None
        }
        with open(manifest_path, 'w', encoding='utf-8') as f:  # type: TextIOWrapper
            json.dump(manifest_entry, f)


class Transcriber:
    """
    A class for transcribing audio files using a pre-trained Whisper model.

    Parameters
    ----------
    model_name : str, optional
        Name of the model to load. Defaults to 'large-v3'.
    device : str, optional
        Device to use for model inference ('cpu' or 'cuda'). Defaults to 'cpu'.
    compute_type : str, optional
        Data type for model computation ('int8', 'float16', etc.). Defaults to 'int8'.

    Attributes
    ----------
    model : faster_whisper.WhisperModel
        Loaded Whisper model for transcription.
    device : str
        Device used for inference.

    Methods
    -------
    transcribe(audio_path, language=None, suppress_numerals=False)
        Transcribes the audio file into text.
    """

    def __init__(
            self,
            model_name: Annotated[str, "Name of the model to load"] = 'large-v3',
            device: Annotated[str, "Device to use for model inference"] = 'cpu',
            compute_type: Annotated[str, "Data type for model computation, e.g., 'int8' or 'float16'"] = 'int8'
    ) -> None:
        if not isinstance(model_name, str):
            raise TypeError("Expected 'model_name' to be of type str")
        if not isinstance(device, str):
            raise TypeError("Expected 'device' to be of type str")
        if not isinstance(compute_type, str):
            raise TypeError("Expected 'compute_type' to be of type str")

        self.device = device
        self.model = faster_whisper.WhisperModel(
            model_name, device=device, compute_type=compute_type
        )

    def transcribe(
            self,
            audio_path: Annotated[str, "Path to the audio file to transcribe"],
            language: Annotated[Optional[str], "Language code for transcription, e.g., 'en' for English"] = None,
            suppress_numerals: Annotated[bool, "Whether to suppress numerals in the transcription"] = False
    ) -> Annotated[Tuple[str, dict], "Transcription text and additional information"]:
        """
        Transcribe an audio file into text.

        Parameters
        ----------
        audio_path : str
            Path to the audio file.
        language : str, optional
            Language code for transcription (e.g., 'en' for English).
        suppress_numerals : bool, optional
            Whether to suppress numerals in the transcription. Defaults to False.

        Returns
        -------
        Tuple[str, dict]
            The transcribed text and additional transcription metadata.

        Examples
        --------
        >>> transcriber = Transcriber()
        >>> text, information = transcriber.transcribe("example.wav")
        >>> isinstance(text, str)
        True
        >>> isinstance(info, dict)
        True
        """
        if not isinstance(audio_path, str):
            raise TypeError("Expected 'audio_path' to be of type str")
        if language is not None and not isinstance(language, str):
            raise TypeError("Expected 'language' to be of type str if provided")
        if not isinstance(suppress_numerals, bool):
            raise TypeError("Expected 'suppress_numerals' to be of type bool")

        audio_waveform = faster_whisper.decode_audio(audio_path)
        suppress_tokens = [-1]
        if suppress_numerals:
            suppress_tokens = TokenizerUtils.find_numeral_symbol_tokens(
                self.model.hf_tokenizer
            )

        transcript_segments, info = self.model.transcribe(
            audio_waveform,
            language=language,
            suppress_tokens=suppress_tokens,
            without_timestamps=True,
            vad_filter=True,
            log_progress=True,
        )

        transcript = ''.join(segment.text for segment in transcript_segments)
        info = vars(info)

        if self.device == 'cuda':
            del self.model
            torch.cuda.empty_cache()

        print(transcript, info)

        return transcript, info


class PunctuationRestorer:
    """
    A class for restoring punctuation in transcribed text.

    Parameters
    ----------
    language : str, optional
        Language for punctuation restoration. Defaults to 'en'.

    Attributes
    ----------
    language : str
        Language used for punctuation restoration.
    punct_model : PunctuationModel
        Model for predicting punctuation.
    supported_languages : List[str]
        List of languages supported by the model.

    Methods
    -------
    restore_punctuation(word_speaker_mapping)
        Restores punctuation in the provided text based on word mappings.
    """

    def __init__(self, language: Annotated[str, "Language for punctuation restoration"] = 'en') -> None:
        self.language = language
        self.punct_model = PunctuationModel(model="kredor/punctuate-all")
        self.supported_languages = [
            "en", "fr", "de", "es", "it", "nl", "pt", "bg", "pl", "cs", "sk", "sl",
        ]

    def restore_punctuation(
            self, word_speaker_mapping: Annotated[List[Dict], "List of word-speaker mappings"]
    ) -> Annotated[List[Dict], "Word mappings with restored punctuation"]:
        """
        Restore punctuation for transcribed text.

        Parameters
        ----------
        word_speaker_mapping : List[Dict]
            List of dictionaries containing word and speaker mappings.

        Returns
        -------
        List[Dict]
            Updated list with punctuation restored.

        Examples
        --------
        >>> restorer = PunctuationRestorer()
        >>> mapping = [{"text": "hello"}, {"text": "world"}]
        >>> result = restorer.restore_punctuation(mapping)
        >>> isinstance(result, list)
        True
        >>> "text" in result[0]
        True
        """
        if self.language not in self.supported_languages:
            print(f"Punctuation restoration is not available for {self.language} language.")
            return word_speaker_mapping

        words_list = [word_dict["text"] for word_dict in word_speaker_mapping]
        labeled_words = self.punct_model.predict(words_list)

        ending_puncts = ".?!"
        model_puncts = ".,;:!?"
        is_acronym = lambda x: re.fullmatch(r"\b(?:[a-zA-Z]\.){2,}", x)

        for word_dict, labeled_tuple in zip(word_speaker_mapping, labeled_words):
            word = word_dict["text"]
            if (
                    word
                    and labeled_tuple[1] in ending_puncts
                    and (word[-1] not in model_puncts or is_acronym(word))
            ):
                word += labeled_tuple[1]
                word = word.rstrip(".") if word.endswith("..") else word
                word_dict["text"] = word

        return word_speaker_mapping


if __name__ == "__main__":
    sample_audio_path = "sample_audio.wav"
    audio_processor_instance = AudioProcessor(sample_audio_path)

    mono_audio_path = audio_processor_instance.convert_to_mono()
    print(f"Mono audio file saved at: {mono_audio_path}")

    audio_duration = audio_processor_instance.get_duration()
    print(f"Audio duration: {audio_duration} seconds")

    converted_audio_path = audio_processor_instance.change_format("mp3")
    print(f"Converted audio file saved at: {converted_audio_path}")

    audio_path_trimmed = audio_processor_instance.trim_audio(0.0, 10.0)
    print(f"Trimmed audio file saved at: {audio_path_trimmed}")

    volume_adjusted_audio_path = audio_processor_instance.adjust_volume(5.0)
    print(f"Volume adjusted audio file saved at: {volume_adjusted_audio_path}")

    additional_audio_path = "additional_audio.wav"
    merged_audio_output_path = audio_processor_instance.merge_audio(additional_audio_path)
    print(f"Merged audio file saved at: {merged_audio_output_path}")

    audio_chunk_paths = audio_processor_instance.split_audio(10.0)
    print(f"Audio chunks saved at: {audio_chunk_paths}")

    output_manifest_path = "output_manifest.json"
    audio_processor_instance.create_manifest(output_manifest_path)
    print(f"Manifest file saved at: {output_manifest_path}")

    transcriber_instance = Transcriber()
    transcribed_text_output, transcription_metadata = transcriber_instance.transcribe(sample_audio_path)
    print(f"Transcribed Text: {transcribed_text_output}")
    print(f"Transcription Info: {transcription_metadata}")

    word_mapping_example = [
        {"text": "hello"},
        {"text": "world"},
        {"text": "this"},
        {"text": "is"},
        {"text": "a"},
        {"text": "test"}
    ]
    punctuation_restorer_instance = PunctuationRestorer()
    punctuation_restored_mapping = punctuation_restorer_instance.restore_punctuation(word_mapping_example)
    print(f"Restored Mapping: {punctuation_restored_mapping}")