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

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

sonusai/utils/asr_manifest_functions/data.py

@@ -1,9 +1 @@
-from dataclasses import dataclass
-
 TranscriptData = dict[str, str]
-
-
-@dataclass(frozen=True)
-class PathInfo:
-    abs_path: str
-    audio_filepath: str

sonusai/utils/asr_manifest_functions/librispeech.py

@@ -1,4 +1,4 @@
-from sonusai.utils.asr_manifest_functions import PathInfo
+from sonusai.utils import PathInfo
 from sonusai.utils.asr_manifest_functions import TranscriptData
 
 

sonusai/utils/asr_manifest_functions/mcgill_speech.py

@@ -1,4 +1,4 @@
-from sonusai.utils.asr_manifest_functions import PathInfo
+from sonusai.utils import PathInfo
 
 
 def get_mcgill_speech_manifest_entry(entry: PathInfo, transcript_data: list[str]) -> dict:

sonusai/utils/asr_manifest_functions/vctk_noisy_speech.py

@@ -1,4 +1,4 @@
-from sonusai.utils.asr_manifest_functions import PathInfo
+from sonusai.utils import PathInfo
 from sonusai.utils.asr_manifest_functions import TranscriptData
 
 

sonusai/utils/braced_glob.py

@@ -1,9 +1,10 @@
 from typing import Generator
+from typing import LiteralString
 from typing import Optional
 from typing import Set
 
 
-def expand_braces(text: str, seen: Optional[Set[str]] = None) -> Generator[str, None, None]:
+def expand_braces(text: LiteralString | str | bytes, seen: Optional[Set[str]] = None) -> Generator[str, None, None]:
     """Brace-expansion pre-processing for glob.
 
     Expand all the braces, then run glob on each of the results.
@@ -16,6 +17,9 @@ def expand_braces(text: str, seen: Optional[Set[str]] = None) -> Generator[str,
     if seen is None:
         seen = set()
 
+    if not isinstance(text, str):
+        text = str(text)
+
     spans = [m.span() for m in re.finditer(r'\{[^{}]*}', text)][::-1]
     alts = [text[start + 1: stop - 1].split(',') for start, stop in spans]
 
@@ -31,7 +35,7 @@ def expand_braces(text: str, seen: Optional[Set[str]] = None) -> Generator[str,
             yield from expand_braces(''.join(replaced), seen)
 
 
-def braced_glob(pathname: str, recursive: bool = False) -> list[str]:
+def braced_glob(pathname: LiteralString | str | bytes, recursive: bool = False) -> list[str]:
     from glob import glob
 
     result = []
@@ -41,7 +45,7 @@ def braced_glob(pathname: str, recursive: bool = False) -> list[str]:
     return result
 
 
-def braced_iglob(pathname: str, recursive: bool = False) -> Generator[str, None, None]:
+def braced_iglob(pathname: LiteralString | str | bytes, recursive: bool = False) -> Generator[str, None, None]:
     from glob import iglob
 
     for expanded_path in expand_braces(pathname):