src/common/paths.py

"""Where things live, and the one rule about images: symlink, never copy.

Each source gets a single symlink at data/unified/images/<source> pointing at
that dataset's root. A manifest image_path is then always

    <source>/<path within that dataset>

so the manifest stays portable and we create four symlinks instead of half a
million.
"""

from __future__ import annotations

import os
from pathlib import Path

DEFAULT_UNIFIED_ROOT = Path("data/unified")


def unified_root(override: str | os.PathLike | None = None) -> Path:
    """Resolve the unified data root: CLI flag, then $UNIFIED_ROOT, then default."""
    if override is not None:
        return Path(override)
    return Path(os.environ.get("UNIFIED_ROOT", DEFAULT_UNIFIED_ROOT))


def images_root(root: Path) -> Path:
    return Path(root) / "images"


def link_source(root: Path, source: str, dataset_root: Path) -> Path:
    """Point data/unified/images/<source> at a dataset root.

    Re-running an extractor re-points the link rather than failing, so moving a
    dataset does not mean rebuilding the manifest.
    """
    dataset_root = Path(dataset_root).resolve()
    if not dataset_root.exists():
        raise FileNotFoundError(f"dataset root does not exist: {dataset_root}")

    link = images_root(root) / source
    link.parent.mkdir(parents=True, exist_ok=True)
    if link.is_symlink():
        if link.readlink() == dataset_root:
            return link
        link.unlink()
    elif link.exists():
        raise FileExistsError(f"{link} exists and is not a symlink; refusing to replace it")

    link.symlink_to(dataset_root, target_is_directory=True)
    return link


def flatten(image_path: str) -> str:
    """'nuscenes/samples/CAM_FRONT/x.jpg' -> 'nuscenes__samples__CAM_FRONT__x'.

    Keeps the whole relative path so names stay unique across sources, and keeps
    the source as a prefix so a stray file is traceable by eye.
    """
    return Path(image_path).with_suffix("").as_posix().replace("/", "__")


def resolve_image(root: Path, image_path: str) -> Path:
    """Manifest image_path -> a real path on disk.

    A local build has the images/<source> symlinks. A root pulled from the Hub
    (src.hub.pull) has none, only the flattened copies under {val,train}/images/,
    so fall back to those.
    """
    local = images_root(root) / image_path
    if local.exists():
        return local
    name = flatten(image_path) + Path(image_path).suffix
    for split in ("val", "train"):
        pulled = Path(root) / split / "images" / name
        if pulled.exists():
            return pulled
    return local


def parts_dir(root: Path) -> Path:
    """Per-source manifest fragments. merge.py concatenates everything here."""
    return Path(root) / "parts"


def part_path(root: Path, source: str) -> Path:
    return parts_dir(root) / f"{source}.parquet"