code/src/lrs_source.py

"""Uniform access to LRS2/LRS3 data, whether it lives in an extracted
directory or inside a .tar shard (LRS3 ships as pretrain.00.tar .. .99.tar,
~1GB each; LRS2 ships as one large tar).

Reading an alignment .txt never touches disk — tarfile hands back the
member's bytes directly from the archive. Reading a video, which cv2 needs
as a real file path, extracts just that one member into a scratch directory;
the caller deletes it immediately after cutting the clips it needs (see
scripts/clip_lrs_words.py). This keeps peak disk usage to a handful of
videos at a time even when the source is a 1GB+ tar shard.
"""
from __future__ import annotations

import os
import posixpath
import tarfile
from dataclasses import dataclass
from typing import Iterator, Optional


@dataclass
class Entry:
    """One paired (txt, mp4) sample, addressable back to its source."""
    source: "LRSSource"
    key: str          # unique id: full relative path, no ext, "/" -> "_"
    txt_member: str   # tar member name, or filesystem path if not a tar
    mp4_member: str
    speaker: str      # parent-directory id (LRS layout: <speaker>/<utt>.mp4)


def _key_and_speaker(rel_path_no_ext: str) -> tuple:
    """Derive (unique key, speaker id) from a relative sample path.

    LRS2/LRS3 lay samples out as <speaker_id>/<utterance_id>, and utterance
    ids repeat across speakers — a basename alone is NOT unique. The key is
    the full relative path with separators flattened; the speaker is the
    immediate parent directory (used for speaker-disjoint splits).
    """
    norm = rel_path_no_ext.replace("\\", "/").strip("/")
    parts = norm.split("/")
    speaker = parts[-2] if len(parts) >= 2 else "unknown"
    return norm.replace("/", "_"), speaker


class LRSSource:
    """One tar shard or one extracted directory tree of LRS2/LRS3 data."""

    def __init__(self, path: str):
        # Random access inside a gzipped tar is pathological (every seek
        # re-decompresses from the byte 0), so .tar.gz gets decompressed to a
        # plain .tar sibling once and we read that instead. ~2x disk for the
        # duration, but per-member reads become O(1) seeks.
        if os.path.isfile(path) and path.endswith((".tar.gz", ".tgz")):
            path = self._ensure_decompressed(path)
        self.path = path
        self.is_tar = os.path.isfile(path) and tarfile.is_tarfile(path)
        self._tar: Optional[tarfile.TarFile] = None
        if self.is_tar:
            self._tar = tarfile.open(path, "r")
        elif not os.path.isdir(path):
            raise FileNotFoundError(f"Not a tar file or directory: {path}")

    @staticmethod
    def _ensure_decompressed(gz_path: str) -> str:
        import gzip
        import shutil
        tar_path = (gz_path[:-3] if gz_path.endswith(".gz")
                    else gz_path[:-4] + ".tar")
        if not os.path.isfile(tar_path):
            print(f"[lrs_source] decompressing {os.path.basename(gz_path)} "
                  f"-> {os.path.basename(tar_path)} (one-time)")
            with gzip.open(gz_path, "rb") as src, open(tar_path, "wb") as dst:
                shutil.copyfileobj(src, dst, length=16 * 1024 * 1024)
        return tar_path

    def close(self) -> None:
        if self._tar is not None:
            self._tar.close()
            self._tar = None

    def __enter__(self) -> "LRSSource":
        return self

    def __exit__(self, *exc) -> None:
        self.close()

    def entries(self) -> Iterator[Entry]:
        """Yield every (txt, mp4) pair found in this source."""
        if self.is_tar:
            names = self._tar.getnames()
            txts = sorted(n for n in names if n.endswith(".txt"))
            mp4s = {n for n in names if n.endswith(".mp4")}
            for txt in txts:
                mp4 = txt[:-4] + ".mp4"
                if mp4 in mp4s:
                    key, speaker = _key_and_speaker(txt[:-4])
                    yield Entry(self, key, txt, mp4, speaker)
        else:
            for root, _, files in os.walk(self.path):
                fileset = set(files)
                for f in sorted(files):
                    if f.endswith(".txt"):
                        mp4 = f[:-4] + ".mp4"
                        if mp4 in fileset:
                            rel = os.path.relpath(
                                os.path.join(root, f[:-4]), self.path)
                            key, speaker = _key_and_speaker(rel)
                            yield Entry(self, key,
                                       os.path.join(root, f),
                                       os.path.join(root, mp4), speaker)

    def read_text(self, member: str) -> str:
        if self.is_tar:
            fh = self._tar.extractfile(member)
            if fh is None:
                return ""
            return fh.read().decode("utf-8", errors="ignore")
        with open(member, encoding="utf-8", errors="ignore") as f:
            return f.read()

    def extract_video(self, member: str, scratch_dir: str) -> str:
        """Return a real filesystem path to the video, extracting if needed.

        The caller owns the returned path when it came from a tar: delete it
        once done (see cleanup_video). Non-tar sources return the original
        path directly and must not be deleted.
        """
        if not self.is_tar:
            return member
        os.makedirs(scratch_dir, exist_ok=True)
        out_path = os.path.join(scratch_dir,
                                os.path.basename(member).replace("/", "_"))
        src = self._tar.extractfile(member)
        with open(out_path, "wb") as dst:
            dst.write(src.read())
        return out_path

    def cleanup_video(self, extracted_path: str) -> None:
        if self.is_tar and os.path.exists(extracted_path):
            os.remove(extracted_path)


def open_sources(*path_specs: str) -> Iterator[LRSSource]:
    """Expand a mix of files/directories/globs into LRSSource objects.

    Each item in path_specs may be a single tar file, a directory, or a glob
    pattern (e.g. "data/lrs3_raw/pretrain.*.tar" to sweep all 100 shards).
    """
    import glob as glob_mod
    for spec in path_specs:
        if not spec:
            continue
        matches = sorted(glob_mod.glob(spec)) if any(
            c in spec for c in "*?[]") else [spec]
        for path in matches:
            yield LRSSource(path)