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

lattifai/caption/formats/__init__.py

@@ -0,0 +1,199 @@
+"""Caption format handlers registry.
+
+This module provides a central registry for all caption format readers and writers.
+Formats are registered using decorators and can be looked up by format ID.
+
+Example:
+    >>> from lattifai.caption.formats import get_reader, get_writer
+    >>> reader = get_reader("srt")
+    >>> supervisions = reader.read("input.srt")
+    >>> writer = get_writer("vtt")
+    >>> writer.write(supervisions, "output.vtt")
+"""
+
+from typing import Dict, List, Optional, Type
+
+from .base import FormatHandler, FormatReader, FormatWriter
+
+# Global registries
+_READERS: Dict[str, Type[FormatReader]] = {}
+_WRITERS: Dict[str, Type[FormatWriter]] = {}
+
+
+def register_reader(format_id: str):
+    """Decorator to register a format reader.
+
+    Args:
+        format_id: Unique identifier for the format (e.g., "srt", "vtt")
+
+    Example:
+        @register_reader("srt")
+        class SRTReader(FormatReader):
+            ...
+    """
+
+    def decorator(cls: Type[FormatReader]) -> Type[FormatReader]:
+        cls.format_id = format_id
+        _READERS[format_id.lower()] = cls
+        return cls
+
+    return decorator
+
+
+def register_writer(format_id: str):
+    """Decorator to register a format writer.
+
+    Args:
+        format_id: Unique identifier for the format
+
+    Example:
+        @register_writer("srt")
+        class SRTWriter(FormatWriter):
+            ...
+    """
+
+    def decorator(cls: Type[FormatWriter]) -> Type[FormatWriter]:
+        cls.format_id = format_id
+        _WRITERS[format_id.lower()] = cls
+        return cls
+
+    return decorator
+
+
+def register_format(format_id: str):
+    """Decorator to register both reader and writer for a format.
+
+    Use this for classes that implement both FormatReader and FormatWriter.
+
+    Args:
+        format_id: Unique identifier for the format
+
+    Example:
+        @register_format("srt")
+        class SRTFormat(FormatHandler):
+            ...
+    """
+
+    def decorator(cls: Type[FormatHandler]) -> Type[FormatHandler]:
+        cls.format_id = format_id
+        _READERS[format_id.lower()] = cls
+        _WRITERS[format_id.lower()] = cls
+        return cls
+
+    return decorator
+
+
+def get_reader(format_id: str) -> Optional[Type[FormatReader]]:
+    """Get a reader class by format ID.
+
+    Args:
+        format_id: Format identifier (case-insensitive)
+
+    Returns:
+        Reader class or None if not found
+    """
+    return _READERS.get(format_id.lower())
+
+
+def get_writer(format_id: str) -> Optional[Type[FormatWriter]]:
+    """Get a writer class by format ID.
+
+    Args:
+        format_id: Format identifier (case-insensitive)
+
+    Returns:
+        Writer class or None if not found
+    """
+    return _WRITERS.get(format_id.lower())
+
+
+def list_readers() -> List[str]:
+    """Get list of all registered reader format IDs."""
+    return sorted(_READERS.keys())
+
+
+def list_writers() -> List[str]:
+    """Get list of all registered writer format IDs."""
+    return sorted(_WRITERS.keys())
+
+
+def detect_format(path: str) -> Optional[str]:
+    """Detect format from file path by checking registered readers.
+
+    Args:
+        path: File path to check
+
+    Returns:
+        Format ID or None if no match found
+    """
+    path_str = str(path)
+
+    # Check if it's content instead of a path
+    is_content = "\n" in path_str or len(path_str) > 500
+
+    # Prioritize specific formats that can detect by content
+    # These often use shared extensions like .vtt, .txt, or .xml
+    priority_formats = ["vtt", "gemini", "premiere_xml"]
+    for format_id in priority_formats:
+        reader_cls = _READERS.get(format_id)
+        if reader_cls and reader_cls.can_read(path_str):
+            return format_id
+
+    if is_content:
+        return None
+
+    # Check each reader's extensions
+    path_lower = path_str.lower()
+    for format_id, reader_cls in _READERS.items():
+        if format_id in priority_formats:
+            continue
+        if reader_cls.can_read(path_lower):
+            return format_id
+
+    # Fallback: try extension directly
+    from pathlib import Path
+
+    try:
+        ext = Path(path_lower).suffix.lstrip(".")
+        if ext in _READERS:
+            return ext
+    except (OSError, ValueError):
+        # Likely content, not a path
+        pass
+
+    return None
+
+
+# Import all format modules to trigger registration
+# Standard formats
+from . import gemini  # YouTube/Gemini markdown
+from . import lrc  # Enhanced LRC with word-level timestamps
+from . import pysubs2  # SRT, ASS, SSA, SUB, SAMI
+from . import sbv  # SubViewer
+from . import tabular  # CSV, TSV, AUD, TXT, JSON
+from . import textgrid  # Praat TextGrid
+from . import ttml  # TTML, IMSC1, EBU-TT-D
+from . import vtt  # WebVTT with YouTube VTT word-level timestamp support
+
+# Professional NLE formats
+from .nle import audition  # Adobe Audition / Pro Tools markers
+from .nle import avid  # Avid DS
+from .nle import fcpxml  # Final Cut Pro XML
+from .nle import premiere  # Adobe Premiere Pro XML
+
+__all__ = [
+    # Base classes
+    "FormatReader",
+    "FormatWriter",
+    "FormatHandler",
+    # Registration
+    "register_reader",
+    "register_writer",
+    "register_format",
+    # Lookup
+    "get_reader",
+    "get_writer",
+    "list_readers",
+    "list_writers",
+    "detect_format",
+]

lattifai/caption/formats/base.py

@@ -0,0 +1,211 @@
+"""Base classes for caption format readers and writers.
+
+This module provides abstract base classes that all format handlers must implement,
+ensuring a consistent interface across different caption formats.
+"""
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+from lhotse.utils import Pathlike
+
+if TYPE_CHECKING:
+    from ..supervision import Supervision
+
+
+class FormatReader(ABC):
+    """Abstract base class for caption format readers.
+
+    All format readers must implement the `read` method to parse caption content
+    and return a list of Supervision objects.
+
+    Class Attributes:
+        format_id: Unique identifier for the format (e.g., "srt", "vtt")
+        extensions: List of file extensions this reader handles (e.g., [".srt"])
+        description: Human-readable description of the format
+    """
+
+    format_id: str = ""
+    extensions: List[str] = []
+    description: str = ""
+
+    @classmethod
+    @abstractmethod
+    def read(
+        cls,
+        source: Union[Pathlike, str],
+        normalize_text: bool = True,
+        **kwargs,
+    ) -> List["Supervision"]:
+        """Read caption content and return list of Supervision objects.
+
+        Args:
+            source: File path or string content
+            normalize_text: Whether to normalize text (strip HTML, etc.)
+            **kwargs: Format-specific options
+
+        Returns:
+            List of Supervision objects with timing and text
+        """
+        pass
+
+    @classmethod
+    def extract_metadata(cls, source: Union[Pathlike, str]) -> Dict[str, str]:
+        """Extract metadata from caption file or content.
+
+        Args:
+            source: File path or string content
+
+        Returns:
+            Dictionary of metadata key-value pairs
+        """
+        return {}
+
+    @classmethod
+    def can_read(cls, path: Union[Pathlike, str]) -> bool:
+        """Check if this reader can handle the given file.
+
+        Args:
+            path: File path to check
+
+        Returns:
+            True if this reader supports the file format
+        """
+        path_str = str(path).lower()
+        return any(path_str.endswith(ext.lower()) for ext in cls.extensions)
+
+    @classmethod
+    def is_content(cls, source: Union[Pathlike, str]) -> bool:
+        """Check if source is string content rather than a file path.
+
+        Args:
+            source: Source to check
+
+        Returns:
+            True if source appears to be content, not a path
+        """
+        if not isinstance(source, str):
+            return False
+        # If it has newlines or is very long, it's likely content
+        return "\n" in source or len(source) > 500
+
+
+class FormatWriter(ABC):
+    """Abstract base class for caption format writers.
+
+    All format writers must implement `write` and `to_bytes` methods.
+
+    Class Attributes:
+        format_id: Unique identifier for the format (e.g., "srt", "vtt")
+        extensions: List of file extensions for this format
+        description: Human-readable description of the format
+    """
+
+    format_id: str = ""
+    extensions: List[str] = []
+    description: str = ""
+
+    @classmethod
+    @abstractmethod
+    def write(
+        cls,
+        supervisions: List["Supervision"],
+        output_path: Pathlike,
+        include_speaker: bool = True,
+        **kwargs,
+    ) -> Path:
+        """Write supervisions to a file.
+
+        Args:
+            supervisions: List of Supervision objects to write
+            output_path: Path to output file
+            include_speaker: Whether to include speaker labels in text
+            **kwargs: Format-specific options
+
+        Returns:
+            Path to the written file
+        """
+        pass
+
+    @classmethod
+    @abstractmethod
+    def to_bytes(
+        cls,
+        supervisions: List["Supervision"],
+        include_speaker: bool = True,
+        **kwargs,
+    ) -> bytes:
+        """Convert supervisions to bytes in this format.
+
+        Args:
+            supervisions: List of Supervision objects
+            include_speaker: Whether to include speaker labels
+            **kwargs: Format-specific options
+
+        Returns:
+            Caption content as bytes
+        """
+        pass
+
+    @classmethod
+    def _should_include_speaker(cls, sup: Any, include_speaker: bool) -> bool:
+        """Check if speaker should be included in output text.
+
+        Considers both the global include_speaker flag and the segment-level
+        'original_speaker' flag in custom metadata.
+        """
+        if not include_speaker or not getattr(sup, "speaker", None):
+            return False
+        custom = getattr(sup, "custom", None)
+        if custom and not custom.get("original_speaker", True):
+            return False
+        return True
+
+
+class FormatHandler(FormatReader, FormatWriter):
+    """Combined reader and writer for formats that support both.
+
+    Most caption formats support both reading and writing. This class
+    combines both interfaces for convenience.
+    """
+
+    pass
+
+
+# Type aliases for registration
+ReaderType = type[FormatReader]
+WriterType = type[FormatWriter]
+
+
+def expand_to_word_supervisions(supervisions: List["Supervision"]) -> List["Supervision"]:
+    """Expand supervisions with word alignment to one supervision per word.
+
+    Used for word-per-segment output when word_level=True but karaoke=False.
+
+    Args:
+        supervisions: List of Supervision objects with optional alignment data
+
+    Returns:
+        List of Supervision objects, one per word if alignment exists,
+        otherwise returns original supervisions unchanged.
+    """
+    from ..supervision import Supervision
+
+    result = []
+    for sup in supervisions:
+        if sup.alignment and "word" in sup.alignment:
+            for word in sup.alignment["word"]:
+                result.append(
+                    Supervision(
+                        text=word.symbol,
+                        start=word.start,
+                        duration=word.duration,
+                        speaker=sup.speaker,
+                        id=f"{sup.id}_word" if sup.id else "",
+                        recording_id=sup.recording_id if hasattr(sup, "recording_id") else "",
+                    )
+                )
+        else:
+            result.append(sup)
+    return result