sonusai 0.17.0__py3-none-any.whl → 0.17.2__py3-none-any.whl

sonusai/speech/l2arctic.py

@@ -0,0 +1,116 @@
+import os
+import string
+from pathlib import Path
+from typing import Optional
+
+from .types import TimeAlignedType
+
+
+def _get_duration(name: str) -> float:
+    import soundfile
+
+    from sonusai import SonusAIError
+
+    try:
+        return soundfile.info(name).duration
+    except Exception as e:
+        raise SonusAIError(f'Error reading {name}: {e}')
+
+
+def load_text(audio: str | os.PathLike[str]) -> Optional[TimeAlignedType]:
+    """Load time-aligned text data given a L2-ARCTIC audio file.
+
+    :param audio: Path to the L2-ARCTIC audio file.
+    :return: A TimeAlignedType object.
+    """
+    file = Path(audio).parent.parent / 'transcript' / (Path(audio).stem + '.txt')
+    if not os.path.exists(file):
+        return None
+
+    with open(file, mode='r', encoding='utf-8') as f:
+        line = f.read()
+
+    return TimeAlignedType(0,
+                           _get_duration(str(audio)),
+                           line.strip().lower().translate(str.maketrans('', '', string.punctuation)))
+
+
+def load_words(audio: str | os.PathLike[str]) -> Optional[list[TimeAlignedType]]:
+    """Load time-aligned word data given a L2-ARCTIC audio file.
+
+    :param audio: Path to the L2-ARCTIC audio file.
+    :return: A list of TimeAlignedType objects.
+    """
+    return _load_ta(audio, 'words')
+
+
+def load_phonemes(audio: str | os.PathLike[str]) -> Optional[list[TimeAlignedType]]:
+    """Load time-aligned phonemes data given a L2-ARCTIC audio file.
+
+    :param audio: Path to the L2-ARCTIC audio file.
+    :return: A list of TimeAlignedType objects.
+    """
+    return _load_ta(audio, 'phones')
+
+
+def _load_ta(audio: str | os.PathLike[str], tier: str) -> Optional[list[TimeAlignedType]]:
+    from praatio import textgrid
+
+    file = Path(audio).parent.parent / 'textgrid' / (Path(audio).stem + '.TextGrid')
+    if not os.path.exists(file):
+        return None
+
+    tg = textgrid.openTextgrid(str(file), includeEmptyIntervals=False)
+    if tier not in tg.tierNames:
+        return None
+
+    entries: list[TimeAlignedType] = []
+    for entry in tg.getTier(tier).entries:
+        entries.append(TimeAlignedType(text=entry.label, start=entry.start, end=entry.end))
+
+    return entries
+
+
+def load_annotations(audio: str | os.PathLike[str]) -> Optional[dict[str, list[TimeAlignedType]]]:
+    """Load time-aligned annotation data given a L2-ARCTIC audio file.
+
+    :param audio: Path to the L2-ARCTIC audio file.
+    :return: A dictionary of a list of TimeAlignedType objects.
+    """
+    from praatio import textgrid
+
+    file = Path(audio).parent.parent / 'annotation' / (Path(audio).stem + '.TextGrid')
+    if not os.path.exists(file):
+        return None
+
+    tg = textgrid.openTextgrid(str(file), includeEmptyIntervals=False)
+    result: dict[str, list[TimeAlignedType]] = {}
+    for tier in tg.tierNames:
+        entries: list[TimeAlignedType] = []
+        for entry in tg.getTier(tier).entries:
+            entries.append(TimeAlignedType(text=entry.label, start=entry.start, end=entry.end))
+        result[tier] = entries
+
+    return result
+
+
+def load_speakers(input_dir: Path) -> dict:
+    speakers = {}
+    with open(input_dir / 'readme-download.txt') as file:
+        processing = False
+        for line in file:
+            if not processing and line.startswith('|---|'):
+                processing = True
+                continue
+
+            if processing:
+                if line.startswith('|**Total**|'):
+                    break
+                else:
+                    fields = line.strip().split('|')
+                    speaker_id = fields[1]
+                    gender = fields[2]
+                    dialect = fields[3]
+                    speakers[speaker_id] = {'gender': gender, 'dialect': dialect}
+
+    return speakers

sonusai/speech/librispeech.py

@@ -0,0 +1,99 @@
+import os
+from pathlib import Path
+from typing import Optional
+
+from .types import TimeAlignedType
+
+
+def _get_num_samples(audio: str | os.PathLike[str]) -> int:
+    """Get number of samples from audio file using soundfile
+
+    :param audio: Audio file name
+    :return: Number of samples
+    """
+    import soundfile
+    from pydub import AudioSegment
+
+    if Path(audio).suffix == '.mp3':
+        return AudioSegment.from_mp3(audio).frame_count()
+
+    if Path(audio).suffix == '.m4a':
+        return AudioSegment.from_file(audio).frame_count()
+
+    return soundfile.info(audio).frames
+
+
+def load_text(audio: str | os.PathLike[str]) -> Optional[TimeAlignedType]:
+    """Load text data from a LibriSpeech transcription file given a LibriSpeech audio filename.
+
+    :param audio: Path to the LibriSpeech audio file.
+    :return: A TimeAlignedType object.
+    """
+    import string
+
+    from sonusai.mixture import get_sample_rate
+
+    path = Path(audio)
+    name = path.stem
+    transcript_filename = path.parent / f'{path.parent.parent.name}-{path.parent.name}.trans.txt'
+
+    if not os.path.exists(transcript_filename):
+        return None
+
+    with open(transcript_filename, mode='r', encoding='utf-8') as f:
+        for line in f.readlines():
+            fields = line.strip().split()
+            key = fields[0]
+            if key == name:
+                text = ' '.join(fields[1:]).lower().translate(str.maketrans('', '', string.punctuation))
+                return TimeAlignedType(0, _get_num_samples(audio) / get_sample_rate(str(audio)), text)
+
+    return None
+
+
+def load_words(audio: str | os.PathLike[str]) -> Optional[list[TimeAlignedType]]:
+    """Load time-aligned word data given a LibriSpeech audio file.
+
+    :param audio: Path to the Librispeech audio file.
+    :return: A list of TimeAlignedType objects.
+    """
+    return _load_ta(audio, 'words')
+
+
+def load_phonemes(audio: str | os.PathLike[str]) -> Optional[list[TimeAlignedType]]:
+    """Load time-aligned phonemes data given a LibriSpeech audio file.
+
+    :param audio: Path to the LibriSpeech audio file.
+    :return: A list of TimeAlignedType objects.
+    """
+    return _load_ta(audio, 'phones')
+
+
+def _load_ta(audio: str | os.PathLike[str], tier: str) -> Optional[list[TimeAlignedType]]:
+    from praatio import textgrid
+
+    file = Path(audio).with_suffix('.TextGrid')
+    if not os.path.exists(file):
+        return None
+
+    tg = textgrid.openTextgrid(str(file), includeEmptyIntervals=False)
+    if tier not in tg.tierNames:
+        return None
+
+    entries: list[TimeAlignedType] = []
+    for entry in tg.getTier(tier).entries:
+        entries.append(TimeAlignedType(text=entry.label, start=entry.start, end=entry.end))
+
+    return entries
+
+
+def load_speakers(input_dir: Path) -> dict:
+    speakers = {}
+    with open(input_dir / 'SPEAKERS.TXT') as file:
+        for line in file:
+            if not line.startswith(';'):
+                fields = line.strip().split('|')
+                speaker_id = fields[0].strip()
+                gender = fields[1].strip()
+                speakers[speaker_id] = {'gender': gender}
+    return speakers

sonusai/speech/mcgill.py

@@ -0,0 +1,70 @@
+import os
+from typing import Optional
+
+from .types import TimeAlignedType
+
+
+def load_text(audio: str | os.PathLike[str]) -> Optional[TimeAlignedType]:
+    """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 sonusai.mixture 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,100 @@
+from pathlib import Path
+
+from praatio import textgrid
+from praatio.utilities.constants import Interval
+
+from .types import TimeAlignedType
+
+
+def _get_duration(name: str) -> float:
+    from pydub import AudioSegment
+
+    from sonusai import SonusAIError
+
+    try:
+        return AudioSegment.from_file(name).duration_seconds
+    except Exception as e:
+        raise SonusAIError(f'Error reading {name}: {e}')
+
+
+def create_textgrid(prompt: Path,
+                    speaker_id: str,
+                    speaker: dict,
+                    output_dir: Path,
+                    text: TimeAlignedType = None,
+                    words: list[TimeAlignedType] = None,
+                    phonemes: list[TimeAlignedType] = None) -> None:
+    if text is not None or words is not None or phonemes is not None:
+        min_t, max_t = _get_min_max({'phonemes': phonemes,
+                                     'text':     [text],
+                                     'words':    words})
+    else:
+        min_t = 0
+        max_t = _get_duration(str(prompt))
+
+    tg = textgrid.Textgrid()
+
+    tg.addTier(textgrid.IntervalTier('speaker_id', [Interval(min_t, max_t, speaker_id)], min_t, max_t))
+    for tier in speaker.keys():
+        tg.addTier(textgrid.IntervalTier(tier, [Interval(min_t, max_t, str(speaker[tier]))], min_t, max_t))
+
+    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, list[TimeAlignedType]]) -> tuple[float, float]:
+    starts = []
+    ends = []
+    for tier in tiers.values():
+        if tier is not None:
+            starts.append(tier[0].start)
+            ends.append(tier[-1].end)
+
+    return min(starts), max(ends)
+
+
+def annotate_textgrid(tiers: dict[str, list[TimeAlignedType]], 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 tier in tiers.keys():
+        entries = []
+        for entry in tiers[tier]:
+            entries.append(Interval(entry.start, entry.end, entry.text))
+        if tier == 'phones':
+            name = 'annotation_phonemes'
+        else:
+            name = 'annotation_' + tier
+        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,135 @@
+import os
+from pathlib import Path
+from typing import Optional
+
+from .types import TimeAlignedType
+
+
+def load_text(audio: str | os.PathLike[str]) -> Optional[TimeAlignedType]:
+    """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 sonusai.mixture 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, mode='r', 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]) -> Optional[list[TimeAlignedType]]:
+    """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]) -> Optional[list[TimeAlignedType]]:
+    """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) -> Optional[list[TimeAlignedType]]:
+    from sonusai.mixture 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, mode='r', 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,52 @@
+import os
+from pathlib import Path
+from typing import Optional
+
+from .types import TimeAlignedType
+
+
+def _get_duration(name: str) -> float:
+    import soundfile
+
+    from sonusai import SonusAIError
+
+    try:
+        return soundfile.info(name).duration
+    except Exception as e:
+        raise SonusAIError(f'Error reading {name}: {e}')
+
+
+def load_text(audio: str | os.PathLike[str]) -> Optional[TimeAlignedType]:
+    """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, mode='r', 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([field for field in fields[3:]])
+                speakers[speaker_id] = {'gender': gender, 'dialect': dialect, 'age': age}
+    return speakers

sonusai/speech/voxceleb2.py

@@ -0,0 +1,86 @@
+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 = {}
+    first = True
+    with open(input_dir / 'vox2_meta_cleansed.csv', newline='') as file:
+        data = csv.reader(file)
+        for row in data:
+            if first:
+                first = False
+            else:
+                speakers[row[0].strip()] = {'gender': row[2].strip(), 'category': row[3].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

@@ -27,11 +27,12 @@ from .max_text_width import max_text_width
 from .model_utils import import_module
 from .numeric_conversion import float_to_int16
 from .numeric_conversion import int16_to_float
-from .onnx_utils import SonusAIMetaData
 from .onnx_utils import add_sonusai_metadata
 from .onnx_utils import get_sonusai_metadata
+from .onnx_utils import load_ort_session
 from .parallel import pp_imap
 from .parallel import pp_tqdm_imap
+from .path_info import PathInfo
 from .print_mixture_details import print_class_count
 from .print_mixture_details import print_mixture_details
 from .ranges import consolidate_range

sonusai/utils/asr_manifest_functions/__init__.py

@@ -1,4 +1,3 @@
-from .data import PathInfo
 from .data import TranscriptData
 from .librispeech import collect_librispeech_transcripts
 from .librispeech import get_librispeech_manifest_entry