lattifai 1.2.1__py3-none-any.whl → 1.2.2__py3-none-any.whl

lattifai/caption/formats/sbv.py

@@ -0,0 +1,147 @@
+"""SubViewer (SBV) format handler.
+
+SBV is YouTube's native subtitle format with the following structure:
+    0:00:00.000,0:00:02.000
+    Text line 1
+
+    0:00:02.000,0:00:04.000
+    Text line 2
+"""
+
+from pathlib import Path
+from typing import List
+
+from ..parsers.text_parser import normalize_text as normalize_text_fn
+from ..parsers.text_parser import parse_speaker_text
+from ..supervision import Supervision
+from . import register_format
+from .base import FormatHandler
+
+
+@register_format("sbv")
+class SBVFormat(FormatHandler):
+    """SubViewer (SBV) format - YouTube's native format."""
+
+    extensions = [".sbv"]
+    description = "SubViewer - YouTube native subtitle format"
+
+    @classmethod
+    def _parse_sbv_timestamp(cls, timestamp: str) -> float:
+        """Parse SBV timestamp (H:MM:SS.mmm) to seconds."""
+        parts = timestamp.strip().split(":")
+        if len(parts) == 3:
+            h, m, s = parts
+            s_parts = s.split(".")
+            seconds = int(h) * 3600 + int(m) * 60 + int(s_parts[0])
+            if len(s_parts) > 1:
+                seconds += int(s_parts[1]) / 1000.0
+            return seconds
+        return 0.0
+
+    @classmethod
+    def _format_sbv_timestamp(cls, seconds: float) -> str:
+        """Format seconds to SBV timestamp (H:MM:SS.mmm)."""
+        h = int(seconds // 3600)
+        m = int((seconds % 3600) // 60)
+        s = int(seconds % 60)
+        ms = int((seconds % 1) * 1000)
+        return f"{h}:{m:02d}:{s:02d}.{ms:03d}"
+
+    @classmethod
+    def read(
+        cls,
+        source,
+        normalize_text: bool = True,
+        **kwargs,
+    ) -> List[Supervision]:
+        """Read SBV format."""
+        # Get content
+        if cls.is_content(source):
+            content = source
+        else:
+            content = Path(source).read_text(encoding="utf-8")
+
+        supervisions = []
+        entries = content.strip().split("\n\n")
+
+        for entry in entries:
+            lines = entry.strip().split("\n")
+            if len(lines) < 2:
+                continue
+
+            # First line: timestamp (H:MM:SS.mmm,H:MM:SS.mmm)
+            timestamp_line = lines[0].strip()
+            text_lines = lines[1:]
+
+            if "," not in timestamp_line:
+                continue
+
+            try:
+                start_str, end_str = timestamp_line.split(",", 1)
+                start = cls._parse_sbv_timestamp(start_str)
+                end = cls._parse_sbv_timestamp(end_str)
+
+                text = " ".join(text_lines).strip()
+                speaker, text = parse_speaker_text(text)
+
+                if normalize_text:
+                    text = normalize_text_fn(text)
+
+                if end > start:
+                    supervisions.append(
+                        Supervision(
+                            text=text,
+                            start=start,
+                            duration=end - start,
+                            speaker=speaker,
+                        )
+                    )
+            except (ValueError, IndexError):
+                continue
+
+        return supervisions
+
+    @classmethod
+    def write(
+        cls,
+        supervisions: List[Supervision],
+        output_path,
+        include_speaker: bool = True,
+        **kwargs,
+    ) -> Path:
+        """Write SBV format."""
+        output_path = Path(output_path)
+        content = cls.to_bytes(supervisions, include_speaker=include_speaker)
+        output_path.write_bytes(content)
+        return output_path
+
+    @classmethod
+    def to_bytes(
+        cls,
+        supervisions: List[Supervision],
+        include_speaker: bool = True,
+        **kwargs,
+    ) -> bytes:
+        """Convert to SBV format bytes."""
+        lines = []
+
+        for i, sup in enumerate(supervisions):
+            start_time = cls._format_sbv_timestamp(sup.start)
+            end_time = cls._format_sbv_timestamp(sup.end)
+            lines.append(f"{start_time},{end_time}")
+
+            text = sup.text.strip() if sup.text else ""
+            if include_speaker and sup.speaker:
+                # Check if speaker should be included
+                include_this_speaker = True
+                if hasattr(sup, "custom") and sup.custom and not sup.custom.get("original_speaker", True):
+                    include_this_speaker = False
+
+                if include_this_speaker:
+                    text = f"{sup.speaker}: {text}"
+            lines.append(text)
+
+            if i < len(supervisions) - 1:
+                lines.append("")
+
+        return "\n".join(lines).encode("utf-8")

lattifai/caption/formats/tabular.py

@@ -0,0 +1,338 @@
+"""Tabular and plain text format handlers.
+
+Handles: CSV, TSV, AUD (Audacity labels), TXT, JSON
+"""
+
+import csv
+import json
+from io import StringIO
+from pathlib import Path
+from typing import List
+
+from ..parsers.text_parser import normalize_text as normalize_text_fn
+from ..parsers.text_parser import parse_speaker_text, parse_timestamp_text
+from ..supervision import Supervision
+from . import register_format
+from .base import FormatHandler
+
+
+@register_format("csv")
+class CSVFormat(FormatHandler):
+    """CSV (Comma-Separated Values) format.
+
+    Format: speaker,start,end,text (with header)
+    Times are in milliseconds.
+    """
+
+    extensions = [".csv"]
+    description = "CSV - tabular subtitle format"
+
+    @classmethod
+    def read(
+        cls,
+        source,
+        normalize_text: bool = True,
+        **kwargs,
+    ) -> List[Supervision]:
+        """Read CSV format."""
+        if cls.is_content(source):
+            lines = list(csv.reader(StringIO(source)))
+        else:
+            with open(source, "r", encoding="utf-8", newline="") as f:
+                lines = list(csv.reader(f))
+
+        if not lines:
+            return []
+
+        # Check for header
+        first_line = [col.strip().lower() for col in lines[0]]
+        has_header = "start" in first_line and "end" in first_line and "text" in first_line
+        has_speaker = "speaker" in first_line
+
+        supervisions = []
+        start_idx = 1 if has_header else 0
+
+        for parts in lines[start_idx:]:
+            if len(parts) < 3:
+                continue
+            try:
+                if has_speaker and len(parts) >= 4:
+                    speaker = parts[0].strip() or None
+                    start = float(parts[1]) / 1000.0
+                    end = float(parts[2]) / 1000.0
+                    text = ",".join(parts[3:]).strip()
+                else:
+                    start = float(parts[0]) / 1000.0
+                    end = float(parts[1]) / 1000.0
+                    text = ",".join(parts[2:]).strip()
+                    speaker = None
+
+                if normalize_text:
+                    text = normalize_text_fn(text)
+
+                if end > start:
+                    supervisions.append(Supervision(text=text, start=start, duration=end - start, speaker=speaker))
+            except (ValueError, IndexError):
+                continue
+
+        return supervisions
+
+    @classmethod
+    def write(cls, supervisions: List[Supervision], output_path, include_speaker: bool = True, **kwargs) -> Path:
+        """Write CSV format."""
+        output_path = Path(output_path)
+        content = cls.to_bytes(supervisions, include_speaker=include_speaker)
+        output_path.write_bytes(content)
+        return output_path
+
+    @classmethod
+    def to_bytes(cls, supervisions: List[Supervision], include_speaker: bool = True, **kwargs) -> bytes:
+        """Convert to CSV format bytes."""
+        output = StringIO()
+        writer = csv.writer(output)
+
+        if include_speaker:
+            writer.writerow(["speaker", "start", "end", "text"])
+            for sup in supervisions:
+                if cls._should_include_speaker(sup, include_speaker):
+                    text = f"{sup.speaker} {sup.text.strip()}"
+                else:
+                    text = sup.text.strip()
+                writer.writerow([sup.speaker or "", round(1000 * sup.start), round(1000 * sup.end), text])
+        else:
+            writer.writerow(["start", "end", "text"])
+            for sup in supervisions:
+                writer.writerow([round(1000 * sup.start), round(1000 * sup.end), sup.text.strip()])
+
+        return output.getvalue().encode("utf-8")
+
+
+@register_format("tsv")
+class TSVFormat(FormatHandler):
+    """TSV (Tab-Separated Values) format.
+
+    Format: speaker\tstart\tend\ttext (with header)
+    Times are in milliseconds.
+    """
+
+    extensions = [".tsv"]
+    description = "TSV - tab-separated subtitle format"
+
+    @classmethod
+    def read(cls, source, normalize_text: bool = True, **kwargs) -> List[Supervision]:
+        """Read TSV format."""
+        if cls.is_content(source):
+            lines = source.strip().split("\n")
+        else:
+            with open(source, "r", encoding="utf-8") as f:
+                lines = f.readlines()
+
+        if not lines:
+            return []
+
+        first_line = lines[0].strip().lower()
+        has_header = "start" in first_line and "end" in first_line and "text" in first_line
+        has_speaker = "speaker" in first_line
+
+        supervisions = []
+        start_idx = 1 if has_header else 0
+
+        for line in lines[start_idx:]:
+            line = line.strip()
+            if not line:
+                continue
+
+            parts = line.split("\t")
+            if len(parts) < 3:
+                continue
+
+            try:
+                if has_speaker and len(parts) >= 4:
+                    speaker = parts[0].strip() or None
+                    start = float(parts[1]) / 1000.0
+                    end = float(parts[2]) / 1000.0
+                    text = "\t".join(parts[3:]).strip()
+                else:
+                    start = float(parts[0]) / 1000.0
+                    end = float(parts[1]) / 1000.0
+                    text = "\t".join(parts[2:]).strip()
+                    speaker = None
+
+                if normalize_text:
+                    text = normalize_text_fn(text)
+
+                if end > start:
+                    supervisions.append(Supervision(text=text, start=start, duration=end - start, speaker=speaker))
+            except (ValueError, IndexError):
+                continue
+
+        return supervisions
+
+    @classmethod
+    def write(cls, supervisions: List[Supervision], output_path, include_speaker: bool = True, **kwargs) -> Path:
+        """Write TSV format."""
+        output_path = Path(output_path)
+        content = cls.to_bytes(supervisions, include_speaker=include_speaker)
+        output_path.write_bytes(content)
+        return output_path
+
+    @classmethod
+    def to_bytes(cls, supervisions: List[Supervision], include_speaker: bool = True, **kwargs) -> bytes:
+        """Convert to TSV format bytes."""
+        lines = []
+        if include_speaker:
+            lines.append("speaker\tstart\tend\ttext")
+            for sup in supervisions:
+                speaker = sup.speaker if cls._should_include_speaker(sup, include_speaker) else ""
+                text = sup.text.strip().replace("\t", " ")
+                lines.append(f"{speaker}\t{round(1000 * sup.start)}\t{round(1000 * sup.end)}\t{text}")
+        else:
+            lines.append("start\tend\ttext")
+            for sup in supervisions:
+                text = sup.text.strip().replace("\t", " ")
+                lines.append(f"{round(1000 * sup.start)}\t{round(1000 * sup.end)}\t{text}")
+
+        return "\n".join(lines).encode("utf-8")
+
+
+@register_format("aud")
+class AUDFormat(FormatHandler):
+    """Audacity Labels format.
+
+    Format: start\tend\t[[speaker]]text
+    Times are in seconds.
+    """
+
+    extensions = [".aud", ".txt"]
+    description = "Audacity Labels format"
+
+    @classmethod
+    def can_read(cls, path) -> bool:
+        """Only handle .aud extension for reading."""
+        return str(path).lower().endswith(".aud")
+
+    @classmethod
+    def read(cls, source, normalize_text: bool = True, **kwargs) -> List[Supervision]:
+        """Read AUD format."""
+        import re
+
+        if cls.is_content(source):
+            lines = source.strip().split("\n")
+        else:
+            with open(source, "r", encoding="utf-8") as f:
+                lines = f.readlines()
+
+        supervisions = []
+        for line in lines:
+            line = line.strip()
+            if not line:
+                continue
+
+            parts = line.split("\t")
+            if len(parts) < 3:
+                continue
+
+            try:
+                start = float(parts[0])
+                end = float(parts[1])
+                text = "\t".join(parts[2:]).strip()
+
+                # Extract speaker from [[speaker]] prefix
+                speaker = None
+                speaker_match = re.match(r"^\[\[([^\]]+)\]\]\s*(.*)$", text)
+                if speaker_match:
+                    speaker = speaker_match.group(1)
+                    text = speaker_match.group(2)
+
+                if normalize_text:
+                    text = normalize_text_fn(text)
+
+                if end > start:
+                    supervisions.append(Supervision(text=text, start=start, duration=end - start, speaker=speaker))
+            except (ValueError, IndexError):
+                continue
+
+        return supervisions
+
+    @classmethod
+    def write(cls, supervisions: List[Supervision], output_path, include_speaker: bool = True, **kwargs) -> Path:
+        """Write AUD format."""
+        output_path = Path(output_path)
+        content = cls.to_bytes(supervisions, include_speaker=include_speaker)
+        output_path.write_bytes(content)
+        return output_path
+
+    @classmethod
+    def to_bytes(cls, supervisions: List[Supervision], include_speaker: bool = True, **kwargs) -> bytes:
+        """Convert to AUD format bytes."""
+        lines = []
+        for sup in supervisions:
+            text = sup.text.strip().replace("\t", " ")
+            if cls._should_include_speaker(sup, include_speaker):
+                text = f"{sup.speaker} {text}"
+            lines.append(f"{sup.start}\t{sup.end}\t{text}")
+
+        return "\n".join(lines).encode("utf-8")
+
+
+@register_format("txt")
+class TXTFormat(FormatHandler):
+    """Plain text format with optional timestamps.
+
+    Format: [start-end] text or [start-end] [speaker]: text
+    """
+
+    extensions = [".txt"]
+    description = "Plain text with optional timestamps"
+
+    @classmethod
+    def read(cls, source, normalize_text: bool = True, **kwargs) -> List[Supervision]:
+        """Read TXT format."""
+        if cls.is_content(source):
+            lines = source.strip().split("\n")
+        else:
+            with open(source, "r", encoding="utf-8") as f:
+                lines = [line.strip() for line in f.readlines()]
+
+        if normalize_text:
+            lines = [normalize_text_fn(line) for line in lines]
+
+        supervisions = []
+        for line in lines:
+            if not line:
+                continue
+
+            start, end, remaining_text = parse_timestamp_text(line)
+            if start is not None and end is not None:
+                speaker, text = parse_speaker_text(remaining_text)
+                supervisions.append(Supervision(text=text, start=start, duration=end - start, speaker=speaker))
+            else:
+                speaker, text = parse_speaker_text(line)
+                supervisions.append(Supervision(text=text, speaker=speaker))
+
+        return supervisions
+
+    @classmethod
+    def write(cls, supervisions: List[Supervision], output_path, include_speaker: bool = True, **kwargs) -> Path:
+        """Write TXT format."""
+        output_path = Path(output_path)
+        content = cls.to_bytes(supervisions, include_speaker=include_speaker)
+        output_path.write_bytes(content)
+        return output_path
+
+    @classmethod
+    def to_bytes(cls, supervisions: List[Supervision], include_speaker: bool = True, **kwargs) -> bytes:
+        """Convert to TXT format bytes."""
+        lines = []
+        for sup in supervisions:
+            text = sup.text or ""
+            if cls._should_include_speaker(sup, include_speaker):
+                text = f"{sup.speaker} {text}"
+            lines.append(f"[{sup.start:.2f}-{sup.end:.2f}] {text}")
+
+        return "\n".join(lines).encode("utf-8")
+
+
+# JSON format moved to json.py for better organization
+# Import here for backwards compatibility
+from .json import JSONFormat  # noqa: F401

lattifai/caption/formats/textgrid.py

@@ -0,0 +1,193 @@
+"""Praat TextGrid format handler.
+
+TextGrid is Praat's native annotation format, commonly used in phonetics research.
+"""
+
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+from lhotse.utils import Pathlike
+
+from ..supervision import Supervision
+from . import register_format
+from .base import FormatHandler
+
+
+@register_format("textgrid")
+class TextGridFormat(FormatHandler):
+    """Praat TextGrid format for phonetic analysis."""
+
+    extensions = [".textgrid"]
+    description = "Praat TextGrid - phonetics research format"
+
+    @classmethod
+    def read(
+        cls,
+        source,
+        normalize_text: bool = True,
+        **kwargs,
+    ) -> List[Supervision]:
+        """Read TextGrid format using tgt library.
+
+        Preserves tier information in Supervision.custom:
+        - textgrid_tier: Original tier name
+        - textgrid_tier_index: Original tier index (for ordering)
+        """
+        from tgt import read_textgrid
+
+        if cls.is_content(source):
+            # Write to temp file for tgt library
+            with tempfile.NamedTemporaryFile(suffix=".textgrid", delete=False, mode="w") as f:
+                f.write(source)
+                temp_path = f.name
+            try:
+                tgt = read_textgrid(temp_path)
+            finally:
+                Path(temp_path).unlink(missing_ok=True)
+        else:
+            tgt = read_textgrid(str(source))
+
+        supervisions = []
+        for tier_idx, tier in enumerate(tgt.tiers):
+            for interval in tier.intervals:
+                supervisions.append(
+                    Supervision(
+                        text=interval.text,
+                        start=interval.start_time,
+                        duration=interval.end_time - interval.start_time,
+                        speaker=tier.name,
+                        custom={
+                            "textgrid_tier": tier.name,
+                            "textgrid_tier_index": tier_idx,
+                        },
+                    )
+                )
+
+        return sorted(supervisions, key=lambda x: x.start)
+
+    @classmethod
+    def write(
+        cls,
+        supervisions: List[Supervision],
+        output_path,
+        include_speaker: bool = True,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ) -> Path:
+        """Write TextGrid format using tgt library.
+
+        Args:
+            supervisions: List of supervisions to write
+            output_path: Output file path
+            include_speaker: Whether to include speaker in text
+            metadata: Optional metadata (for API consistency)
+        """
+        from lhotse.supervision import AlignmentItem
+        from tgt import Interval, IntervalTier, TextGrid, write_to_file
+
+        output_path = Path(output_path)
+        tg = TextGrid()
+
+        utterances = []
+        words = []
+        scores = {"utterances": [], "words": []}
+
+        for sup in sorted(supervisions, key=lambda x: x.start):
+            text = sup.text or ""
+            if include_speaker and sup.speaker:
+                # Check if speaker should be included
+                include_this_speaker = True
+                if hasattr(sup, "custom") and sup.custom and not sup.custom.get("original_speaker", True):
+                    include_this_speaker = False
+
+                if include_this_speaker:
+                    text = f"{sup.speaker} {text}"
+
+            utterances.append(Interval(sup.start, sup.end, text))
+
+            # Extract word-level alignment if present
+            alignment = getattr(sup, "alignment", None)
+            if alignment and "word" in alignment:
+                for item in alignment["word"]:
+                    words.append(Interval(item.start, item.end, item.symbol))
+                    if item.score is not None:
+                        scores["words"].append(Interval(item.start, item.end, f"{item.score:.2f}"))
+
+            if hasattr(sup, "custom") and sup.custom and "score" in sup.custom:
+                scores["utterances"].append(Interval(sup.start, sup.end, f"{sup.custom['score']:.2f}"))
+
+        tg.add_tier(IntervalTier(name="utterances", objects=utterances))
+
+        if words:
+            tg.add_tier(IntervalTier(name="words", objects=words))
+
+        if scores["utterances"]:
+            tg.add_tier(IntervalTier(name="utterance_scores", objects=scores["utterances"]))
+        if scores["words"]:
+            tg.add_tier(IntervalTier(name="word_scores", objects=scores["words"]))
+
+        write_to_file(tg, str(output_path), format="long")
+        return output_path
+
+    @classmethod
+    def to_bytes(
+        cls,
+        supervisions: List[Supervision],
+        include_speaker: bool = True,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ) -> bytes:
+        """Convert to TextGrid format bytes.
+
+        Args:
+            supervisions: List of supervisions to convert
+            include_speaker: Whether to include speaker in text
+            metadata: Optional metadata (currently unused, for API consistency)
+        """
+        # TextGrid requires file I/O due to tgt library implementation
+        with tempfile.NamedTemporaryFile(suffix=".textgrid", delete=False) as tmp:
+            tmp_path = Path(tmp.name)
+
+        try:
+            cls.write(supervisions, tmp_path, include_speaker, metadata=metadata, **kwargs)
+            return tmp_path.read_bytes()
+        finally:
+            tmp_path.unlink(missing_ok=True)
+
+    @classmethod
+    def extract_metadata(cls, source: Union[Pathlike, str], **kwargs) -> Dict[str, Any]:
+        """Extract metadata from TextGrid.
+
+        Returns:
+            Dict containing:
+            - textgrid_xmin: Minimum time boundary
+            - textgrid_xmax: Maximum time boundary
+            - textgrid_tiers: List of tier names
+        """
+        import re
+        from pathlib import Path
+
+        metadata: Dict[str, Any] = {}
+        if cls.is_content(source):
+            content = source
+        else:
+            try:
+                with open(source, "r", encoding="utf-8") as f:
+                    content = f.read()
+            except Exception:
+                return {}
+
+        match = re.search(r"xmin\s*=\s*([\d.]+)", content)
+        if match:
+            metadata["textgrid_xmin"] = float(match.group(1))
+        match = re.search(r"xmax\s*=\s*([\d.]+)", content)
+        if match:
+            metadata["textgrid_xmax"] = float(match.group(1))
+
+        # Extract tier names
+        tier_names = re.findall(r'name\s*=\s*"([^"]+)"', content)
+        if tier_names:
+            metadata["textgrid_tiers"] = tier_names
+
+        return metadata