depthdif_dataset/dataloaders.py

from __future__ import annotations

from typing import Any

import torch
from torch.utils.data import DataLoader, Dataset, Subset, random_split


def build_dataloader(
    dataset: Dataset,
    *,
    batch_size: int = 16,
    shuffle: bool = True,
    num_workers: int = 0,
    pin_memory: bool = True,
    persistent_workers: bool = False,
    prefetch_factor: int | None = 2,
) -> DataLoader:
    """Build a PyTorch DataLoader with the DepthDif training defaults."""
    num_workers = int(num_workers)
    kwargs: dict[str, Any] = dict(
        dataset=dataset,
        batch_size=int(batch_size),
        shuffle=bool(shuffle),
        num_workers=num_workers,
        pin_memory=bool(pin_memory),
        persistent_workers=bool(persistent_workers) and num_workers > 0,
    )
    # PyTorch only accepts prefetch_factor when worker processes are enabled.
    if num_workers > 0 and prefetch_factor is not None:
        kwargs["prefetch_factor"] = int(prefetch_factor)
    return DataLoader(**kwargs)


def split_dataset(
    dataset: Dataset,
    *,
    val_fraction: float = 0.2,
    seed: int = 7,
) -> tuple[Subset, Subset]:
    """Create a deterministic train/validation split from one dataset."""
    total_len = len(dataset)
    if total_len == 0:
        raise RuntimeError("Dataset is empty; cannot create train/val split.")

    val_len = int(round(total_len * float(val_fraction)))
    if total_len > 1:
        val_len = min(max(val_len, 1 if val_fraction > 0.0 else 0), total_len - 1)
    else:
        val_len = 0
    train_len = total_len - val_len
    generator = torch.Generator().manual_seed(int(seed))
    train_dataset, val_dataset = random_split(
        dataset,
        [train_len, val_len],
        generator=generator,
    )
    return train_dataset, val_dataset


def build_train_val_dataloaders(
    dataset: Dataset,
    *,
    val_dataset: Dataset | None = None,
    dataloader_cfg: dict[str, Any] | None = None,
    val_fraction: float = 0.2,
    seed: int = 7,
) -> tuple[DataLoader, DataLoader]:
    """Build train and validation DataLoaders from one or two datasets."""
    cfg = dict(dataloader_cfg or {})
    if val_dataset is None:
        train_dataset, val_dataset = split_dataset(
            dataset,
            val_fraction=val_fraction,
            seed=seed,
        )
    else:
        train_dataset = dataset

    train_loader = build_dataloader(
        train_dataset,
        batch_size=int(cfg.get("batch_size", 16)),
        shuffle=bool(cfg.get("shuffle", True)),
        num_workers=int(cfg.get("num_workers", 4)),
        pin_memory=bool(cfg.get("pin_memory", True)),
        persistent_workers=bool(cfg.get("persistent_workers", False)),
        prefetch_factor=cfg.get("prefetch_factor", 2),
    )
    val_loader = build_dataloader(
        val_dataset,
        batch_size=int(cfg.get("val_batch_size", cfg.get("batch_size", 16))),
        # Keep the repository's intended behavior: validation is shuffled by default.
        shuffle=bool(cfg.get("val_shuffle", True)),
        num_workers=int(cfg.get("val_num_workers", 0)),
        pin_memory=bool(cfg.get("pin_memory", True)),
        persistent_workers=bool(cfg.get("val_persistent_workers", False)),
        prefetch_factor=cfg.get("prefetch_factor", 2),
    )
    return train_loader, val_loader


class DepthTileDataModule:
    """Small PyTorch-only DataModule-style wrapper for DepthDif tiles."""

    def __init__(
        self,
        *,
        dataset: Dataset,
        val_dataset: Dataset | None = None,
        dataloader_cfg: dict[str, Any] | None = None,
        val_fraction: float = 0.2,
        seed: int = 7,
    ) -> None:
        """Store dataset and loader settings without requiring Lightning."""
        self.dataset = dataset
        self.val_dataset = val_dataset
        self.dataloader_cfg = dataloader_cfg or {}
        self.val_fraction = float(val_fraction)
        self.seed = int(seed)
        self.train_dataset: Subset | Dataset | None = (
            dataset if val_dataset is not None else None
        )
        self._train_val_split_done = val_dataset is not None

    def setup(self, stage: str | None = None) -> None:
        """Prepare deterministic train/validation datasets."""
        _ = stage
        if self._train_val_split_done:
            return
        self.train_dataset, self.val_dataset = split_dataset(
            self.dataset,
            val_fraction=self.val_fraction,
            seed=self.seed,
        )
        self._train_val_split_done = True

    def train_dataloader(self) -> DataLoader:
        """Return the configured training DataLoader."""
        if not self._train_val_split_done:
            self.setup("fit")
        return build_train_val_dataloaders(
            self.train_dataset,
            val_dataset=self.val_dataset,
            dataloader_cfg=self.dataloader_cfg,
            val_fraction=self.val_fraction,
            seed=self.seed,
        )[0]

    def val_dataloader(self) -> DataLoader:
        """Return the configured validation DataLoader."""
        if not self._train_val_split_done:
            self.setup("fit")
        return build_train_val_dataloaders(
            self.train_dataset,
            val_dataset=self.val_dataset,
            dataloader_cfg=self.dataloader_cfg,
            val_fraction=self.val_fraction,
            seed=self.seed,
        )[1]