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

lattifai/caption/utils.py

@@ -0,0 +1,474 @@
+"""Utility functions for caption processing.
+
+This module provides utility functions for:
+- Timecode offset handling (for professional timelines starting at 01:00:00:00)
+- Overlap/collision resolution (merge or trim modes)
+- SRT format optimization (UTF-8 BOM, comma-separated milliseconds)
+"""
+
+from copy import deepcopy
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING, List, Optional, Tuple
+
+if TYPE_CHECKING:
+    from .supervision import Supervision
+
+
+class CollisionMode(Enum):
+    """Mode for resolving overlapping captions."""
+
+    MERGE = "merge"  # Merge overlapping lines with line break
+    TRIM = "trim"  # Trim earlier caption to end before later starts
+    KEEP = "keep"  # Keep overlaps as-is (may cause issues in some NLE)
+
+
+@dataclass
+class TimecodeOffset:
+    """Configuration for timecode offset.
+
+    Professional timelines often start at 01:00:00:00 instead of 00:00:00:00.
+    This class handles the offset conversion.
+
+    Attributes:
+        hours: Hour offset (default 0)
+        minutes: Minute offset (default 0)
+        seconds: Second offset (default 0)
+        frames: Frame offset (default 0)
+        fps: Frame rate for frame-based offset calculation
+    """
+
+    hours: int = 0
+    minutes: int = 0
+    seconds: float = 0.0
+    frames: int = 0
+    fps: float = 25.0
+
+    @property
+    def total_seconds(self) -> float:
+        """Calculate total offset in seconds."""
+        return self.hours * 3600 + self.minutes * 60 + self.seconds + (self.frames / self.fps)
+
+    @classmethod
+    def from_timecode(cls, timecode: str, fps: float = 25.0) -> "TimecodeOffset":
+        """Create offset from timecode string.
+
+        Args:
+            timecode: Timecode string (HH:MM:SS:FF or HH:MM:SS.mmm)
+            fps: Frame rate
+
+        Returns:
+            TimecodeOffset instance
+        """
+        # Handle different separators
+        if ";" in timecode:
+            # Drop-frame format
+            parts = timecode.replace(";", ":").split(":")
+        else:
+            parts = timecode.split(":")
+
+        hours = int(parts[0]) if len(parts) > 0 else 0
+        minutes = int(parts[1]) if len(parts) > 1 else 0
+
+        # Handle seconds (may have frames or milliseconds)
+        if len(parts) > 2:
+            sec_part = parts[2]
+            if "." in sec_part:
+                # Millisecond format
+                seconds = float(sec_part)
+                frames = 0
+            else:
+                seconds = float(sec_part)
+                frames = int(parts[3]) if len(parts) > 3 else 0
+        else:
+            seconds = 0.0
+            frames = 0
+
+        return cls(hours=hours, minutes=minutes, seconds=seconds, frames=frames, fps=fps)
+
+    @classmethod
+    def broadcast_start(cls, fps: float = 25.0) -> "TimecodeOffset":
+        """Create standard broadcast start offset (01:00:00:00).
+
+        Args:
+            fps: Frame rate
+
+        Returns:
+            TimecodeOffset for broadcast start
+        """
+        return cls(hours=1, fps=fps)
+
+
+def apply_timecode_offset(
+    supervisions: List["Supervision"],
+    offset: TimecodeOffset,
+) -> List["Supervision"]:
+    """Apply timecode offset to all supervisions.
+
+    Args:
+        supervisions: List of supervision segments
+        offset: Timecode offset to apply
+
+    Returns:
+        New list of supervisions with offset applied
+    """
+    from .supervision import Supervision
+
+    offset_seconds = offset.total_seconds
+    result = []
+
+    for sup in supervisions:
+        new_sup = Supervision(
+            text=sup.text,
+            start=sup.start + offset_seconds,
+            duration=sup.duration,
+            speaker=sup.speaker,
+            id=sup.id,
+            language=sup.language,
+            alignment=deepcopy(getattr(sup, "alignment", None)),
+            custom=sup.custom.copy() if sup.custom else None,
+        )
+
+        # Also offset word-level alignments if present
+        if new_sup.alignment and "word" in new_sup.alignment:
+            from lhotse.supervision import AlignmentItem
+
+            new_words = []
+            for word in new_sup.alignment["word"]:
+                new_words.append(
+                    AlignmentItem(
+                        symbol=word.symbol,
+                        start=word.start + offset_seconds,
+                        duration=word.duration,
+                        score=word.score,
+                    )
+                )
+            new_sup.alignment["word"] = new_words
+
+        result.append(new_sup)
+
+    return result
+
+
+def resolve_overlaps(
+    supervisions: List["Supervision"],
+    mode: CollisionMode = CollisionMode.MERGE,
+    gap_threshold: float = 0.05,
+) -> List["Supervision"]:
+    """Resolve overlapping supervisions.
+
+    Args:
+        supervisions: List of supervision segments (should be sorted by start time)
+        mode: How to handle overlaps (MERGE, TRIM, or KEEP)
+        gap_threshold: Minimum gap between captions in seconds (for TRIM mode)
+
+    Returns:
+        New list of supervisions with overlaps resolved
+    """
+    from .supervision import Supervision
+
+    if not supervisions or mode == CollisionMode.KEEP:
+        return supervisions
+
+    # Sort by start time
+    sorted_sups = sorted(supervisions, key=lambda x: x.start)
+    result = []
+
+    i = 0
+    while i < len(sorted_sups):
+        current = sorted_sups[i]
+
+        # Find all overlapping supervisions
+        overlapping = [current]
+        j = i + 1
+        while j < len(sorted_sups):
+            next_sup = sorted_sups[j]
+            # Check if next overlaps with any in our group
+            current_end = max(s.end for s in overlapping)
+            if next_sup.start < current_end:
+                overlapping.append(next_sup)
+                j += 1
+            else:
+                break
+
+        if len(overlapping) == 1:
+            # No overlap
+            result.append(current)
+            i += 1
+        elif mode == CollisionMode.MERGE:
+            # Merge all overlapping into one
+            merged = _merge_supervisions(overlapping)
+            result.append(merged)
+            i = j
+        elif mode == CollisionMode.TRIM:
+            # Trim each to not overlap with next
+            for k, sup in enumerate(overlapping[:-1]):
+                next_sup = overlapping[k + 1]
+                # Trim current to end before next starts
+                new_duration = max(gap_threshold, next_sup.start - sup.start - gap_threshold)
+                trimmed = Supervision(
+                    text=sup.text,
+                    start=sup.start,
+                    duration=min(sup.duration, new_duration),
+                    speaker=sup.speaker,
+                    id=sup.id,
+                    language=sup.language,
+                    alignment=sup.alignment,
+                    custom=sup.custom,
+                )
+                result.append(trimmed)
+            # Add last one as-is
+            result.append(overlapping[-1])
+            i = j
+        else:
+            result.append(current)
+            i += 1
+
+    return result
+
+
+def _merge_supervisions(supervisions: List["Supervision"]) -> "Supervision":
+    """Merge multiple overlapping supervisions into one.
+
+    Args:
+        supervisions: List of overlapping supervisions
+
+    Returns:
+        Single merged supervision
+    """
+    from .supervision import Supervision
+
+    if not supervisions:
+        raise ValueError("Cannot merge empty supervision list")
+
+    if len(supervisions) == 1:
+        return supervisions[0]
+
+    # Calculate merged timing
+    start = min(s.start for s in supervisions)
+    end = max(s.end for s in supervisions)
+
+    # Merge text with line breaks, indicating speakers
+    texts = []
+    for sup in supervisions:
+        text = sup.text.strip() if sup.text else ""
+        if sup.speaker:
+            texts.append(f"- {sup.speaker}: {text}")
+        else:
+            texts.append(f"- {text}")
+
+    merged_text = "\n".join(texts)
+
+    # Use first supervision's speaker or None for mixed speakers
+    speakers = set(s.speaker for s in supervisions if s.speaker)
+    speaker = supervisions[0].speaker if len(speakers) == 1 else None
+
+    return Supervision(
+        text=merged_text,
+        start=start,
+        duration=end - start,
+        speaker=speaker,
+        id=supervisions[0].id,
+        language=supervisions[0].language,
+    )
+
+
+def format_srt_timestamp(seconds: float) -> str:
+    """Format timestamp for SRT format.
+
+    SRT uses comma as millisecond separator: HH:MM:SS,mmm
+
+    Args:
+        seconds: Time in seconds
+
+    Returns:
+        SRT-formatted timestamp string
+    """
+    if seconds < 0:
+        seconds = 0
+
+    hours = int(seconds // 3600)
+    minutes = int((seconds % 3600) // 60)
+    secs = int(seconds % 60)
+    millis = int((seconds % 1) * 1000)
+
+    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
+
+
+def generate_srt_content(
+    supervisions: List["Supervision"],
+    include_speaker: bool = True,
+    use_bom: bool = True,
+) -> bytes:
+    """Generate SRT content with proper formatting.
+
+    Args:
+        supervisions: List of supervision segments
+        include_speaker: Include speaker labels in text
+        use_bom: Include UTF-8 BOM for Windows compatibility
+
+    Returns:
+        SRT content as bytes
+    """
+    lines = []
+
+    for i, sup in enumerate(supervisions, 1):
+        # Sequence number
+        lines.append(str(i))
+
+        # Timestamp line with comma separator
+        start_ts = format_srt_timestamp(sup.start)
+        end_ts = format_srt_timestamp(sup.end)
+        lines.append(f"{start_ts} --> {end_ts}")
+
+        # Text content
+        text = sup.text.strip() if sup.text else ""
+        if include_speaker and sup.speaker:
+            # Check if speaker was originally in text
+            if not (hasattr(sup, "custom") and sup.custom and not sup.custom.get("original_speaker", True)):
+                text = f"{sup.speaker}: {text}"
+        lines.append(text)
+
+        # Blank line between entries
+        lines.append("")
+
+    content = "\n".join(lines)
+
+    if use_bom:
+        # UTF-8 with BOM for Windows compatibility
+        return b"\xef\xbb\xbf" + content.encode("utf-8")
+    else:
+        return content.encode("utf-8")
+
+
+def detect_overlaps(supervisions: List["Supervision"]) -> List[Tuple[int, int]]:
+    """Detect all overlapping supervision pairs.
+
+    Args:
+        supervisions: List of supervision segments
+
+    Returns:
+        List of tuples (index1, index2) where supervisions overlap
+    """
+    overlaps = []
+    sorted_sups = sorted(enumerate(supervisions), key=lambda x: x[1].start)
+
+    for i in range(len(sorted_sups) - 1):
+        idx1, sup1 = sorted_sups[i]
+        for j in range(i + 1, len(sorted_sups)):
+            idx2, sup2 = sorted_sups[j]
+            if sup2.start >= sup1.end:
+                break
+            overlaps.append((idx1, idx2))
+
+    return overlaps
+
+
+def split_long_lines(
+    supervisions: List["Supervision"],
+    max_chars_per_line: int = 42,
+    max_lines: int = 2,
+) -> List["Supervision"]:
+    """Split supervisions with long text into multiple segments.
+
+    Useful for broadcast compliance where line length limits are strict.
+
+    Args:
+        supervisions: List of supervision segments
+        max_chars_per_line: Maximum characters per line
+        max_lines: Maximum lines per supervision
+
+    Returns:
+        New list with long supervisions split
+    """
+    from .supervision import Supervision
+
+    result = []
+    max_total_chars = max_chars_per_line * max_lines
+
+    for sup in supervisions:
+        text = sup.text.strip() if sup.text else ""
+
+        if len(text) <= max_total_chars:
+            # Text fits, just wrap lines if needed
+            wrapped = _wrap_text(text, max_chars_per_line, max_lines)
+            new_sup = Supervision(
+                text=wrapped,
+                start=sup.start,
+                duration=sup.duration,
+                speaker=sup.speaker,
+                id=sup.id,
+                language=sup.language,
+                alignment=sup.alignment,
+                custom=sup.custom,
+            )
+            result.append(new_sup)
+        else:
+            # Split into multiple supervisions
+            chunks = _split_text_chunks(text, max_total_chars)
+            chunk_duration = sup.duration / len(chunks)
+
+            for i, chunk in enumerate(chunks):
+                wrapped = _wrap_text(chunk, max_chars_per_line, max_lines)
+                new_sup = Supervision(
+                    text=wrapped,
+                    start=sup.start + i * chunk_duration,
+                    duration=chunk_duration,
+                    speaker=sup.speaker if i == 0 else None,
+                    id=f"{sup.id}_{i}" if sup.id else None,
+                    language=sup.language,
+                )
+                result.append(new_sup)
+
+    return result
+
+
+def _wrap_text(text: str, max_chars: int, max_lines: int) -> str:
+    """Wrap text to fit within character and line limits."""
+    words = text.split()
+    lines = []
+    current_line = []
+    current_length = 0
+
+    for word in words:
+        word_len = len(word)
+        if current_length + word_len + (1 if current_line else 0) <= max_chars:
+            current_line.append(word)
+            current_length += word_len + (1 if len(current_line) > 1 else 0)
+        else:
+            if current_line:
+                lines.append(" ".join(current_line))
+            current_line = [word]
+            current_length = word_len
+
+            if len(lines) >= max_lines:
+                break
+
+    if current_line and len(lines) < max_lines:
+        lines.append(" ".join(current_line))
+
+    return "\n".join(lines[:max_lines])
+
+
+def _split_text_chunks(text: str, max_chars: int) -> List[str]:
+    """Split text into chunks that fit within character limit."""
+    words = text.split()
+    chunks = []
+    current_chunk = []
+    current_length = 0
+
+    for word in words:
+        word_len = len(word)
+        if current_length + word_len + (1 if current_chunk else 0) <= max_chars:
+            current_chunk.append(word)
+            current_length += word_len + (1 if len(current_chunk) > 1 else 0)
+        else:
+            if current_chunk:
+                chunks.append(" ".join(current_chunk))
+            current_chunk = [word]
+            current_length = word_len
+
+    if current_chunk:
+        chunks.append(" ".join(current_chunk))
+
+    return chunks

lattifai/cli/__init__.py

@@ -4,7 +4,7 @@ import nemo_run as run  # noqa: F401
 
 # Import and re-export entrypoints at package level so NeMo Run can find them
 from lattifai.cli.alignment import align
-from lattifai.cli.caption import convert
+from lattifai.cli.caption import convert, diff
 from lattifai.cli.diarization import diarize
 from lattifai.cli.transcribe import transcribe, transcribe_align
 from lattifai.cli.youtube import youtube
@@ -12,6 +12,7 @@ from lattifai.cli.youtube import youtube
 __all__ = [
     "align",
     "convert",
+    "diff",
     "diarize",
     "transcribe",
     "transcribe_align",

lattifai/cli/caption.py

@@ -7,6 +7,7 @@ from lhotse.utils import Pathlike
 from typing_extensions import Annotated
 
 from lattifai.config import CaptionConfig
+from lattifai.config.caption import KaraokeConfig
 from lattifai.utils import safe_print
 
 
@@ -16,6 +17,8 @@ def convert(
     output_path: Pathlike,
     include_speaker_in_text: bool = False,
     normalize_text: bool = False,
+    word_level: bool = False,
+    karaoke: bool = False,
 ):
     """
     Convert caption file to another format.
@@ -33,6 +36,11 @@ def convert(
         normalize_text: Whether to normalize caption text during conversion.
             This applies text cleaning such as removing HTML tags, decoding entities,
             collapsing whitespace, and standardizing punctuation.
+        word_level: Use word-level output format if supported.
+            When True without karaoke: outputs word-per-segment (each word as separate segment).
+            JSON format will include a 'words' field with word-level timestamps.
+        karaoke: Enable karaoke styling (requires word_level=True).
+            When True: outputs karaoke format (ASS \\kf tags, enhanced LRC, etc.).
 
     Examples:
         # Basic format conversion (positional arguments)
@@ -41,6 +49,15 @@ def convert(
         # Convert with text normalization
         lai caption convert input.srt output.json normalize_text=true
 
+        # Convert to word-per-segment output (if input has alignment)
+        lai caption convert input.json output.srt word_level=true
+
+        # Convert to karaoke format (ASS with \\kf tags)
+        lai caption convert input.json output.ass word_level=true karaoke=true
+
+        # Export JSON with word-level timestamps
+        lai caption convert input.srt output.json word_level=true
+
         # Mixing positional and keyword arguments
         lai caption convert input.srt output.vtt \\
             include_speaker_in_text=false \\
@@ -53,8 +70,16 @@ def convert(
     """
     from lattifai.caption import Caption
 
+    # Create karaoke_config if karaoke flag is set
+    karaoke_config = KaraokeConfig(enabled=True) if karaoke else None
+
     caption = Caption.read(input_path, normalize_text=normalize_text)
-    caption.write(output_path, include_speaker_in_text=include_speaker_in_text)
+    caption.write(
+        output_path,
+        include_speaker_in_text=include_speaker_in_text,
+        word_level=word_level,
+        karaoke_config=karaoke_config,
+    )
 
     safe_print(f"✅ Converted {input_path} -> {output_path}")
     return output_path
@@ -178,6 +203,88 @@ def shift(
     return output_path
 
 
+@run.cli.entrypoint(name="diff", namespace="caption")
+def diff(
+    ref_path: Pathlike,
+    hyp_path: Pathlike,
+    split_sentence: bool = True,
+    verbose: bool = True,
+):
+    """
+    Compare and align caption supervisions with transcription segments.
+
+    This command reads a reference caption file and a hypothesis file, then performs
+    text alignment to show how they match up. It's useful for comparing
+    original subtitles against ASR (Automatic Speech Recognition) results.
+
+    Args:
+        ref_path: Path to reference caption file (ground truth)
+        hyp_path: Path to hypothesis file (e.g., ASR results)
+        split_sentence: Enable sentence splitting before alignment (default: True)
+        verbose: Enable verbose output to show detailed alignment info (default: True)
+
+    Examples:
+        # Compare reference with hypothesis (positional arguments)
+        lai caption diff subtitles.srt transcription.json
+
+        # Disable sentence splitting
+        lai caption diff subtitles.srt transcription.json split_sentence=false
+
+        # Disable verbose output
+        lai caption diff subtitles.srt transcription.json verbose=false
+    """
+    from pathlib import Path
+
+    from lattifai.alignment.sentence_splitter import SentenceSplitter
+    from lattifai.alignment.text_align import align_supervisions_and_transcription
+    from lattifai.caption import Caption
+
+    ref_path = Path(ref_path).expanduser()
+    hyp_path = Path(hyp_path).expanduser()
+
+    # Read reference caption (supervisions)
+    caption_obj = Caption.read(ref_path)
+
+    # Read hypothesis
+    hyp_obj = Caption.read(hyp_path)
+
+    # Apply sentence splitting if enabled
+    if split_sentence:
+        splitter = SentenceSplitter(device="cpu", lazy_init=True)
+        caption_obj.supervisions = splitter.split_sentences(caption_obj.supervisions)
+        hyp_obj.supervisions = splitter.split_sentences(hyp_obj.supervisions)
+
+    # Set transcription on caption object
+    caption_obj.transcription = hyp_obj.supervisions
+
+    safe_print(f"📖  Reference: {len(caption_obj.supervisions)} segments from {ref_path}")
+    safe_print(f"🎤 Hypothesis: {len(caption_obj.transcription)} segments from {hyp_path}")
+    if split_sentence:
+        safe_print("✂️  Sentence splitting: enabled")
+    safe_print("")
+
+    # Perform alignment
+    results = align_supervisions_and_transcription(
+        caption=caption_obj,
+        verbose=verbose,
+    )
+
+    # # Print summary
+    # safe_print("")
+    # safe_print("=" * 72)
+    # safe_print(f"📊 Alignment Summary: {len(results)} groups")
+    # for idx, (sub_align, asr_align, quality, timestamp, typing) in enumerate(results):
+    #     sub_count = len(sub_align) if sub_align else 0
+    #     asr_count = len(asr_align) if asr_align else 0
+    #     safe_print(f"  Group {idx + 1}: ref={sub_count}, hyp={asr_count}, {quality.info}, typing={typing}")
+
+    return results
+
+
+def main_diff():
+    run.cli.main(diff)
+
+
 def main_convert():
     run.cli.main(convert)
 

lattifai/cli/transcribe.py

@@ -108,12 +108,7 @@ def transcribe(
     is_url = media_config.is_input_remote()
 
     # Prepare output paths
-    if is_url:
-        # For URLs, use output_dir from media_config or current directory
-        output_path = media_config.output_dir
-    else:
-        # For files, use input path directory
-        output_path = Path(media_config.input_path).parent
+    output_dir = media_config.output_dir or Path(media_config.input_path).parent
 
     # Create transcriber
     if not transcription_config.lattice_model_path:
@@ -134,13 +129,13 @@ def transcribe(
         if is_url:
             # Download media first, then transcribe
             safe_print(colorful.cyan("    Downloading media from URL..."))
-            from lattifai.workflow.youtube import YouTubeDownloader
+            from lattifai.youtube import YouTubeDownloader
 
             downloader = YouTubeDownloader()
             input_path = asyncio.run(
                 downloader.download_media(
                     url=media_config.input_path,
-                    output_dir=str(output_path),
+                    output_dir=str(output_dir),
                     media_format=media_config.normalize_format(),
                     force_overwrite=media_config.force_overwrite,
                 )
@@ -167,7 +162,7 @@ def transcribe(
         if is_url:
             # For URLs, generate output filename based on transcriber
             output_format = transcriber.file_suffix.lstrip(".")
-            final_output = output_path / f"youtube_LattifAI_{transcriber.name}.{output_format}"
+            final_output = output_dir / f"youtube_LattifAI_{transcriber.name}.{output_format}"
         else:
             # For files, use input filename with suffix
             final_output = Path(media_config.input_path).with_suffix(".LattifAI.srt")

lattifai/cli/youtube.py

@@ -44,7 +44,8 @@ def youtube(
     Args:
         yt_url: YouTube video URL (can be provided as positional argument)
         media: Media configuration for controlling formats and output directories.
-            Fields: input_path (YouTube URL), output_dir, output_format, force_overwrite
+            Fields: input_path (YouTube URL), output_dir, output_format, force_overwrite,
+                    audio_track_id (default: "original"), quality (default: "best")
         client: API client configuration.
             Fields: api_key, timeout, max_retries
         alignment: Alignment configuration (model selection and inference settings).
@@ -129,6 +130,8 @@ def youtube(
         channel_selector=media_config.channel_selector,
         streaming_chunk_secs=media_config.streaming_chunk_secs,
         use_transcription=use_transcription,
+        audio_track_id=media_config.audio_track_id,
+        quality=media_config.quality,
     )