code/src/dataset.py

"""Unified dataset that merges any number of preprocessed clip roots
(e.g. LRS2-derived clips, LRS3-derived clips, your own recordings).

Every source uses the same on-disk layout, produced by scripts/prepare_data.py:

    <root>/<WORD>/<split>/<clip_id>.npy      # uint8 (T, H, W), 0..255

Because every source shares this layout, they combine transparently:
build_datasets() scans every root in cfg.data_roots, unions the word list
into a single label map, and returns train/val/test datasets drawn from all
of them.
"""
from __future__ import annotations

import glob
import json
import os
from typing import List, Tuple

import numpy as np
import torch
from torch.utils.data import Dataset

from src.transforms import ClipTransform

SPLITS = ("train", "val", "test")


def _scan_words(roots: List[str]) -> List[str]:
    words = set()
    for root in roots:
        if not root or not os.path.isdir(root):
            continue
        for name in os.listdir(root):
            if os.path.isdir(os.path.join(root, name)):
                words.add(name)
    return sorted(words)


def _collect(roots: List[str], split: str,
             label_map: dict) -> List[Tuple[str, int]]:
    items: List[Tuple[str, int]] = []
    for root in roots:
        if not root or not os.path.isdir(root):
            continue
        for word, label in label_map.items():
            pattern = os.path.join(root, word, split, "*.npy")
            for path in glob.glob(pattern):
                items.append((path, label))
    return items


class ClipDataset(Dataset):
    def __init__(self, items: List[Tuple[str, int]], transform: ClipTransform):
        self.items = items
        self.transform = transform

    def __len__(self) -> int:
        return len(self.items)

    def __getitem__(self, idx: int):
        path, label = self.items[idx]
        clip = np.load(path)                 # (T, H, W) uint8 or float32
        if clip.dtype == np.uint8:
            clip = clip.astype(np.float32) / 255.0
        return self.transform(clip), label


class DummyDataset(Dataset):
    """Synthetic clips for smoke-testing the pipeline end to end."""

    def __init__(self, n: int, num_classes: int, num_frames: int, size: int):
        self.n = n
        self.num_classes = num_classes
        self.num_frames = num_frames
        self.size = size

    def __len__(self) -> int:
        return self.n

    def __getitem__(self, idx: int):
        label = idx % self.num_classes
        rng = np.random.default_rng(idx)
        # Give each class a distinct brightness bias so the model can learn.
        clip = rng.random((1, self.num_frames, self.size, self.size),
                          dtype=np.float32) * 0.2 + label / self.num_classes
        return torch.from_numpy(clip), label


def build_datasets(cfg):
    """Return (train_ds, val_ds, test_ds, label_map)."""
    if cfg.dummy:
        label_map = {f"word{i}": i for i in range(cfg.dummy_classes)}
        make = lambda n: DummyDataset(n, cfg.dummy_classes,
                                      cfg.num_frames, cfg.image_size)
        return (make(cfg.dummy_samples), make(cfg.dummy_samples // 4),
                make(cfg.dummy_samples // 4), label_map)

    roots = [r.strip() for r in cfg.data_roots.split(",") if r.strip()]
    words = ([w.strip() for w in cfg.words.split(",") if w.strip()]
             or _scan_words(roots))
    if not words:
        raise RuntimeError(
            "No word folders found. Run scripts/prepare_data.py first, "
            "point --data_roots at your data (comma-separated), or use --dummy.")
    label_map = {w: i for i, w in enumerate(words)}

    train_tf = ClipTransform(cfg.image_size, cfg.num_frames, train=True)
    eval_tf = ClipTransform(cfg.image_size, cfg.num_frames, train=False)

    train_ds = ClipDataset(_collect(roots, "train", label_map), train_tf)
    val_ds = ClipDataset(_collect(roots, "val", label_map), eval_tf)
    test_ds = ClipDataset(_collect(roots, "test", label_map), eval_tf)
    return train_ds, val_ds, test_ds, label_map


def save_label_map(label_map: dict, out_dir: str) -> None:
    os.makedirs(out_dir, exist_ok=True)
    with open(os.path.join(out_dir, "label_map.json"), "w") as f:
        json.dump(label_map, f, indent=2)