sonusai 1.0.16__cp311-abi3-macosx_11_0_arm64.whl

sonusai/speech/mcgill.py

@@ -0,0 +1,71 @@
+import os
+
+from .types import TimeAlignedType
+
+
+def load_text(audio: str | os.PathLike[str]) -> TimeAlignedType | None:
+    """Load time-aligned text data given a McGill-Speech audio file.
+
+    :param audio: Path to the McGill-Speech audio file.
+    :return: A TimeAlignedType object.
+    """
+    import string
+    import struct
+
+    from ..mixture.audio import get_sample_rate
+
+    if not os.path.exists(audio):
+        return None
+
+    sample_rate = get_sample_rate(str(audio))
+
+    with open(audio, mode="rb") as f:
+        content = f.read()
+
+    riff_id, file_size, wave_id = struct.unpack("<4si4s", content[:12])
+    if riff_id.decode("utf-8") != "RIFF":
+        return None
+
+    if wave_id.decode("utf-8") != "WAVE":
+        return None
+
+    fmt_id, fmt_size = struct.unpack("<4si", content[12:20])
+
+    if fmt_id.decode("utf-8") != "fmt ":
+        return None
+
+    if fmt_size != 16:
+        return None
+
+    (
+        _wave_format_tag,
+        channels,
+        _samples_per_sec,
+        _avg_bytes_per_sec,
+        _block_align,
+        bits_per_sample,
+    ) = struct.unpack("<hhiihh", content[20:36])
+
+    i = 36
+    samples = None
+    text = None
+    while i < file_size:
+        chunk_id = struct.unpack("<4s", content[i : i + 4])[0].decode("utf-8")
+        chunk_size = struct.unpack("<i", content[i + 4 : i + 8])[0]
+
+        if chunk_id == "data":
+            samples = chunk_size / channels / (bits_per_sample / 8)
+            break
+
+        if chunk_id == "afsp":
+            chunks = struct.unpack(f"<{chunk_size}s", content[i + 8 : i + 8 + chunk_size])[0]
+            chunks = chunks.decode("utf-8").split("\x00")
+            for chunk in chunks:
+                if chunk.startswith('text: "'):
+                    text = chunk[7:-1].lower().translate(str.maketrans("", "", string.punctuation))
+        i += 8 + chunk_size + chunk_size % 2
+
+    if text and samples:
+        return TimeAlignedType(start=0, end=samples / sample_rate, text=text)
+
+    return None

sonusai/speech/textgrid.py

@@ -0,0 +1,89 @@
+from pathlib import Path
+
+from praatio import textgrid
+from praatio.utilities.constants import Interval
+
+from .types import TimeAlignedType
+
+
+def create_textgrid(
+    prompt: Path,
+    output_dir: Path,
+    text: TimeAlignedType | None = None,
+    words: list[TimeAlignedType] | None = None,
+    phonemes: list[TimeAlignedType] | None = None,
+) -> None:
+    if text is None and words is None and phonemes is None:
+        return
+
+    min_t, max_t = _get_min_max({"phonemes": phonemes, "text": text, "words": words})
+
+    tg = textgrid.Textgrid()
+
+    if text is not None:
+        entries = [Interval(text.start, text.end, text.text)]
+        text_tier = textgrid.IntervalTier("text", entries, min_t, max_t)
+        tg.addTier(text_tier)
+
+    if words is not None:
+        entries = []
+        for word in words:
+            entries.append(Interval(word.start, word.end, word.text))
+        words_tier = textgrid.IntervalTier("words", entries, min_t, max_t)
+        tg.addTier(words_tier)
+
+    if phonemes is not None:
+        entries = []
+        for phoneme in phonemes:
+            entries.append(Interval(phoneme.start, phoneme.end, phoneme.text))
+        phonemes_tier = textgrid.IntervalTier("phonemes", entries, min_t, max_t)
+        tg.addTier(phonemes_tier)
+
+    output_filename = str(output_dir / prompt.stem) + ".TextGrid"
+    tg.save(output_filename, format="long_textgrid", includeBlankSpaces=True)
+
+
+def _get_min_max(tiers: dict[str, TimeAlignedType | list[TimeAlignedType] | None]) -> tuple[float, float]:
+    starts = []
+    ends = []
+    for tier in tiers.values():
+        if tier is None:
+            continue
+        if isinstance(tier, TimeAlignedType):
+            starts.append(tier.start)
+            ends.append(tier.end)
+        else:
+            starts.append(tier[0].start)
+            ends.append(tier[-1].end)
+
+    return min(starts), max(ends)
+
+
+def annotate_textgrid(
+    tiers: dict[str, TimeAlignedType | list[TimeAlignedType] | None] | None, prompt: Path, output_dir: Path
+) -> None:
+    import os
+
+    if tiers is None:
+        return
+
+    file = Path(output_dir / prompt.stem).with_suffix(".TextGrid")
+    if not os.path.exists(file):
+        tg = textgrid.Textgrid()
+        min_t, max_t = _get_min_max(tiers)
+    else:
+        tg = textgrid.openTextgrid(str(file), includeEmptyIntervals=False)
+        min_t = tg.minTimestamp
+        max_t = tg.maxTimestamp
+
+    for k, v in tiers.items():
+        if v is None:
+            continue
+        entries = [Interval(entry.start, entry.end, entry.text) for entry in v]
+        if k == "phones":
+            name = "annotation_phonemes"
+        else:
+            name = "annotation_" + k
+        tg.addTier(textgrid.IntervalTier(name, entries, min_t, max_t))
+
+    tg.save(str(file), format="long_textgrid", includeBlankSpaces=True)

sonusai/speech/timit.py

@@ -0,0 +1,138 @@
+import os
+from pathlib import Path
+
+from .types import TimeAlignedType
+
+
+def load_text(audio: str | os.PathLike[str]) -> TimeAlignedType | None:
+    """Load time-aligned text data given a TIMIT audio file.
+
+    :param audio: Path to the TIMIT audio file.
+    :return: A TimeAlignedType object.
+    """
+    import string
+
+    from ..mixture.audio import get_sample_rate
+
+    file = Path(audio).with_suffix(".TXT")
+    if not os.path.exists(file):
+        return None
+
+    sample_rate = get_sample_rate(str(audio))
+
+    with open(file, encoding="utf-8") as f:
+        line = f.read()
+
+    fields = line.strip().split()
+    start = int(fields[0]) / sample_rate
+    end = int(fields[1]) / sample_rate
+    text = " ".join(fields[2:]).lower().translate(str.maketrans("", "", string.punctuation))
+
+    return TimeAlignedType(start, end, text)
+
+
+def load_words(audio: str | os.PathLike[str]) -> list[TimeAlignedType] | None:
+    """Load time-aligned word data given a TIMIT audio file.
+
+    :param audio: Path to the TIMIT audio file.
+    :return: A list of TimeAlignedType objects.
+    """
+
+    return _load_ta(audio, "words")
+
+
+def load_phonemes(audio: str | os.PathLike[str]) -> list[TimeAlignedType] | None:
+    """Load time-aligned phonemes data given a TIMIT audio file.
+
+    :param audio: Path to the TIMIT audio file.
+    :return: A list of TimeAlignedType objects.
+    """
+
+    return _load_ta(audio, "phonemes")
+
+
+def _load_ta(audio: str | os.PathLike[str], tier: str) -> list[TimeAlignedType] | None:
+    from ..mixture.audio import get_sample_rate
+
+    if tier == "words":
+        file = Path(audio).with_suffix(".WRD")
+    elif tier == "phonemes":
+        file = Path(audio).with_suffix(".PHN")
+    else:
+        raise ValueError(f"Unknown tier: {tier}")
+
+    if not os.path.exists(file):
+        return None
+
+    sample_rate = get_sample_rate(str(audio))
+
+    entries: list[TimeAlignedType] = []
+    first = True
+    with open(file, encoding="utf-8") as f:
+        for line in f.readlines():
+            fields = line.strip().split()
+            start = int(fields[0]) / sample_rate
+            end = int(fields[1]) / sample_rate
+            text = " ".join(fields[2:])
+
+            if first:
+                first = False
+            else:
+                if start < entries[-1].end:
+                    start = entries[-1].end - (entries[-1].end - start) // 2
+                    entries[-1] = TimeAlignedType(text=entries[-1].text, start=entries[-1].start, end=start)
+
+                if end <= start:
+                    end = start + 1 / sample_rate
+
+            entries.append(TimeAlignedType(text=text, start=start, end=end))
+
+    return entries
+
+
+def _years_between(record, born):
+    try:
+        rec_fields = [int(x) for x in record.split("/")]
+        brn_fields = [int(x) for x in born.split("/")]
+        return rec_fields[2] - brn_fields[2] - ((rec_fields[1], rec_fields[0]) < (brn_fields[1], brn_fields[0]))
+    except ValueError:
+        return "??"
+
+
+def _decode_dialect(d: str) -> str:
+    if d in ["DR1", "1"]:
+        return "New England"
+    if d in ["DR2", "2"]:
+        return "Northern"
+    if d in ["DR3", "3"]:
+        return "North Midland"
+    if d in ["DR4", "4"]:
+        return "South Midland"
+    if d in ["DR5", "5"]:
+        return "Southern"
+    if d in ["DR6", "6"]:
+        return "New York City"
+    if d in ["DR7", "7"]:
+        return "Western"
+    if d in ["DR8", "8"]:
+        return "Army Brat"
+
+    raise ValueError(f"Unrecognized dialect: {d}")
+
+
+def load_speakers(input_dir: Path) -> dict:
+    speakers = {}
+    with open(input_dir / "SPKRINFO.TXT") as file:
+        for line in file:
+            if not line.startswith(";"):
+                fields = line.strip().split()
+                speaker_id = fields[0]
+                gender = fields[1]
+                dialect = _decode_dialect(fields[2])
+                age = _years_between(fields[4], fields[5])
+                speakers[speaker_id] = {
+                    "gender": gender,
+                    "dialect": dialect,
+                    "age": age,
+                }
+    return speakers

sonusai/speech/types.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class TimeAlignedType:
+    start: float
+    end: float
+    text: str
+
+    @property
+    def duration(self) -> float:
+        return self.end - self.start

sonusai/speech/vctk.py

@@ -0,0 +1,53 @@
+import os
+from pathlib import Path
+
+from .types import TimeAlignedType
+
+
+def _get_duration(name: str) -> float:
+    import soundfile
+
+    try:
+        return soundfile.info(name).duration
+    except Exception as e:
+        raise OSError(f"Error reading {name}: {e}") from e
+
+
+def load_text(audio: str | os.PathLike[str]) -> TimeAlignedType | None:
+    """Load time-aligned text data given a VCTK audio file.
+
+    :param audio: Path to the VCTK audio file.
+    :return: A TimeAlignedType object.
+    """
+    import string
+
+    file = Path(audio).parents[2] / "txt" / Path(audio).parent.name / (Path(audio).stem[:-5] + ".txt")
+    if not os.path.exists(file):
+        return None
+
+    with open(file, encoding="utf-8") as f:
+        line = f.read()
+
+    start = 0
+    end = _get_duration(str(audio))
+    text = line.strip().lower().translate(str.maketrans("", "", string.punctuation))
+
+    return TimeAlignedType(start, end, text)
+
+
+def load_speakers(input_dir: Path) -> dict:
+    speakers = {}
+    with open(input_dir / "speaker-info.txt") as file:
+        for line in file:
+            if not line.startswith("ID"):
+                fields = line.strip().split("(", 1)[0].split()
+                speaker_id = fields[0]
+                age = fields[1]
+                gender = fields[2]
+                dialect = " ".join(list(fields[3:]))
+                speakers[speaker_id] = {
+                    "gender": gender,
+                    "dialect": dialect,
+                    "age": age,
+                }
+    return speakers

sonusai/speech/voxceleb.py

@@ -0,0 +1,108 @@
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass(frozen=True)
+class Segment:
+    person: str
+    video: str
+    id: str
+    start: float
+    stop: float
+
+
+def load_speakers(input_dir: Path) -> dict:
+    import csv
+
+    speakers = {}
+
+    # VoxCeleb1
+    first = True
+    with open(input_dir / "vox1_meta.csv", newline="") as file:
+        data = csv.reader(file, delimiter="\t")
+        for row in data:
+            if first:
+                first = False
+            else:
+                speakers[row[0].strip()] = {
+                    "gender": row[2].strip(),
+                    "dialect": row[3].strip(),
+                    "category": row[4].strip(),
+                }
+
+    # VoxCeleb2
+    first = True
+    with open(input_dir / "vox2_meta.csv", newline="") as file:
+        data = csv.reader(file, delimiter="\t")
+        for row in data:
+            if first:
+                first = False
+            else:
+                speakers[row[1].strip()] = {
+                    "gender": row[3].strip(),
+                    "category": row[4].strip(),
+                }
+
+    return speakers
+
+
+def load_segment(path: str | os.PathLike[str]) -> Segment:
+    path = Path(path)
+
+    with path.open("r") as file:
+        segment = file.read().strip()
+
+    header, frames = segment.split("\n\n")
+    header_dict = _parse_header(header)
+    start, stop = _get_segment_boundaries(frames)
+
+    return Segment(
+        person=header_dict["Identity"],
+        video=header_dict["Reference"],
+        id=path.stem,
+        start=start,
+        stop=stop,
+    )
+
+
+def _parse_header(header: str) -> dict:
+    def _parse_line(line: str) -> tuple[str, str]:
+        """Parse a line of header text into a dictionary.
+
+        Header text has the following format:
+
+        Identity  : \tid00017
+        Reference : \t7t6lfzvVaTM
+        Offset    : \t1
+        FV Conf   : \t16.647\t(1)
+        ASD Conf  : \t4.465
+
+        """
+        k, v = line.split("\t", maxsplit=1)
+        k = k[:-2].strip()
+        v = v.strip()
+        return k, v
+
+    return dict(_parse_line(line) for line in header.split("\n"))
+
+
+def _get_segment_boundaries(frames: str) -> tuple[float, float]:
+    """Get the start and stop points of the segment.
+
+    Frames text has the following format:
+
+    FRAME 	X 	Y 	W 	H
+    000245 	0.392 	0.223 	0.253 	0.451
+    ...
+    000470 	0.359 	0.207 	0.260 	0.463
+
+    """
+
+    def _get_frame_seconds(line: str) -> float:
+        frame = int(line.split("\t")[0])
+        # YouTube is 25 FPS
+        return frame / 25
+
+    lines = frames.split("\n")
+    return _get_frame_seconds(lines[1]), _get_frame_seconds(lines[-1])

sonusai/utils/__init__.py

@@ -0,0 +1,3 @@
+from .asr import ASRResult
+
+__all__ = ["ASRResult"]

sonusai/utils/asl_p56.py

@@ -0,0 +1,130 @@
+from ..datatypes import AudioT
+
+
+def asl_p56(audio: AudioT) -> float:
+    """Implement ITU-T P.56 method B
+    :param audio: audio for which to calculate active speech level
+    :return: Active speech level mean square energy
+    """
+    import numpy as np
+    import scipy.signal as signal
+
+    from ..constants import SAMPLE_RATE
+
+    eps = np.finfo(np.float32).eps
+
+    # Time constant of smoothing in seconds
+    T = 0.03
+
+    # Coefficient of smoothing
+    g = np.exp(-1 / (SAMPLE_RATE * T))
+
+    # Hangover time in seconds
+    H = 0.2
+    # Rounded up to next integer
+    H_samples = np.ceil(H * SAMPLE_RATE)
+
+    # Margin in dB, difference between threshold and active speech level
+    M = 15.9
+
+    # Number of thresholds
+    thresh_num = 15
+
+    # Series of fixed threshold voltages to apply to the envelope. These are spaced
+    # in geometric progression, at intervals of not more than 2:1 (6.02 dB), from a
+    # value equal to about half the maximum code down to a value equal to one
+    # quantizing interval or lower.
+    c = 2 ** np.arange(-15, thresh_num - 15, dtype=np.float32)
+
+    # Activity counter for each threshold
+    a = np.full(thresh_num, -1)
+
+    # Hangover counter for each threshold
+    h = np.full(thresh_num, H_samples)
+
+    # Long-term level square energy of audio
+    sq = sum(np.square(audio))
+
+    # Use a 2nd order IIR filter to detect the envelope q
+    p = signal.lfilter([1 - g, 0], [1, -g], abs(audio))
+    # q is the envelope, obtained from moving average of abs(audio) (with slight "hangover").
+    q = signal.lfilter([1 - g, 0], [1, -g], p)
+
+    for k in range(len(audio)):
+        for j in range(thresh_num):
+            if q[k] >= c[j]:
+                a[j] = a[j] + 1
+                h[j] = 0
+            elif h[j] < H_samples:
+                a[j] = a[j] + 1
+                h[j] = h[j] + 1
+            else:
+                break
+    asl_msq = 0
+    if a[0] == -1:
+        return asl_msq
+
+    a += 2
+    A_db1 = 10 * np.log10(sq / a[0] + eps)
+    C_db1 = 20 * np.log10(c[0] + eps)
+    if A_db1 - C_db1 < M:
+        return asl_msq
+
+    A_db = np.zeros(thresh_num)
+    C_db = np.zeros(thresh_num)
+    delta = np.zeros(thresh_num)
+    A_db[0] = A_db1
+    C_db[0] = C_db1
+    delta[0] = A_db1 - C_db1
+
+    for j in range(1, thresh_num):
+        A_db[j] = 10 * np.log10(sq / (a[j] + eps) + eps)
+        C_db[j] = 20 * np.log10(c[j] + eps)
+
+    for j in range(1, thresh_num):
+        if a[j] != 0:
+            delta[j] = A_db[j] - C_db[j]
+            if delta[j] <= M:
+                # Interpolate to find the asl_ms_log
+                asl_ms_log = _bin_interp(A_db[j], A_db[j - 1], C_db[j], C_db[j - 1], M, 0.5)
+                # This is the mean square value NOT the RMS
+                asl_msq = 10.0 ** (asl_ms_log / 10)
+                break
+
+    return asl_msq
+
+
+def _bin_interp(u_cnt: float, l_cnt: float, u_thr: float, l_thr: float, margin: float, tol: float) -> float:
+    tol = abs(tol)
+
+    # Check if extreme counts are not already the true active value
+    iter_num = 1
+    if abs(u_cnt - u_thr - margin) < tol:
+        return u_cnt
+
+    if abs(l_cnt - l_thr - margin) < tol:
+        return l_cnt
+
+    # Initialize first middle for given (initial) bounds
+    m_cnt = (u_cnt + l_cnt) / 2.0
+    m_thr = (u_thr + l_thr) / 2.0
+
+    while True:
+        # Loop until diff falls inside the tolerance (-tol<=diff<=tol)
+        diff = m_cnt - m_thr - margin
+        if abs(diff) <= tol:
+            break
+
+        # If tolerance is not met up to 20 iterations, then relax the tolerance by 10
+        iter_num += 1
+        if iter_num > 20:
+            tol = tol * 1.1
+
+        if diff > tol:
+            m_cnt = (u_cnt + m_cnt) / 2.0
+            m_thr = (u_thr + m_thr) / 2.0
+        elif diff < -tol:
+            m_cnt = (m_cnt + l_cnt) / 2.0
+            m_thr = (m_thr + l_thr) / 2.0
+
+    return m_cnt

sonusai/utils/asr.py

@@ -0,0 +1,91 @@
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from ..datatypes import AudioT
+
+
+@dataclass(frozen=True)
+class ASRResult:
+    text: str
+    confidence: float | None = None
+    lang: str | None = None
+    lang_prob: float | None = None
+    duration: float | None = None
+    num_segments: int | None = None
+    asr_cpu_time: float | None = None
+
+
+def get_available_engines() -> list[str]:
+    from importlib import import_module
+    from pkgutil import iter_modules
+
+    module = import_module("sonusai.utils.asr_functions")
+    engines = [method for method in dir(module) if not method.startswith("_")]
+    for _, name, _ in iter_modules():
+        if name.startswith("sonusai_asr_"):
+            module = import_module(f"{name}.asr_functions")
+            for method in dir(module):
+                if not method.startswith("_"):
+                    engines.append(method)
+
+    return engines
+
+
+def _asr_fn(engine: str) -> Callable[..., ASRResult]:
+    from importlib import import_module
+    from pkgutil import iter_modules
+
+    module = import_module("sonusai.utils.asr_functions")
+    for method in dir(module):
+        if method == engine:
+            return getattr(module, method)
+
+    for _, name, _ in iter_modules():
+        if name.startswith("sonusai_asr_"):
+            module = import_module(f"{name}.asr_functions")
+            for method in dir(module):
+                if method == engine:
+                    return getattr(module, method)
+
+    raise ValueError(f"engine {engine} not supported")
+
+
+def calc_asr(audio: AudioT | str, engine: str, **config) -> ASRResult:
+    """Run ASR on audio
+
+    :param audio: Numpy array of audio samples or location of an audio file
+    :param engine: ASR engine to use
+    :param config: kwargs configuration parameters
+    :return: ASRResult object containing text and confidence
+    """
+    from copy import copy
+
+    import numpy as np
+
+    from ..mixture.audio import read_audio
+
+    if not isinstance(audio, np.ndarray):
+        audio = copy(read_audio(audio, config.get("use_cache", True)))
+
+    return _asr_fn(engine)(audio, **config)
+
+
+def validate_asr(engine: str, **config) -> None:
+    from importlib import import_module
+    from pkgutil import iter_modules
+
+    module = import_module("sonusai.utils.asr_functions")
+    for method in dir(module):
+        if method == engine:
+            getattr(module, method + "_validate")(**config)
+            return
+
+    for _, name, _ in iter_modules():
+        if name.startswith("sonusai_asr_"):
+            module = import_module(f"{name}.asr_functions")
+            for method in dir(module):
+                if method == engine:
+                    getattr(module, method + "_validate")(**config)
+                    return
+
+    raise ValueError(f"engine {engine} not supported")

sonusai/utils/asr_functions/__init__.py

@@ -0,0 +1,3 @@
+from .aaware_whisper import aaware_whisper
+
+__all__ = ["aaware_whisper"]