dataforge: robocap end-to-end slice (download/convert/register/view)

packages/dataforge/dataforge/apis/__init__.py

@@ -0,0 +1 @@
+"""Verb entry points: download, convert, register, view (one module each)."""

packages/dataforge/dataforge/apis/convert.py

@@ -0,0 +1,46 @@
+"""``dataforge-convert``: write the base-layer rrd for one or every sequence."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from dataforge.datasets import AnnotatedDatasetUnion, RobocapConfig
+from dataforge.datasets.base import DataforgeDataset, DataforgeDatasetConfig
+from dataforge.identity import SequenceIdentity
+
+
+@dataclass
+class Config:
+    """Convert discovered sequences into base-layer recordings."""
+
+    dataset: AnnotatedDatasetUnion = field(default_factory=RobocapConfig)
+    """Dataset to convert; the raw-tree location lives on the dataset config."""
+    sequence: str | None = None
+    """Convert only this sequence (its ``sequence_key``, ``recording_id``, or any part)."""
+    force: bool = False
+    """Rewrite recordings that already exist instead of skipping them."""
+
+
+def select(identities: list[SequenceIdentity], sequence: str | None) -> list[SequenceIdentity]:
+    """Filter discovered identities by key, recording id, or a single part."""
+    if sequence is None:
+        return identities
+    matches: list[SequenceIdentity] = [
+        identity for identity in identities if sequence in (identity.sequence_key, identity.recording_id) or sequence in identity.parts
+    ]
+    if not matches:
+        raise ValueError(f"No sequence matches {sequence!r} among {len(identities)} discovered sequences")
+    return matches
+
+
+def main(config: Config) -> None:
+    """Convert every selected sequence serially."""
+    dataset_config: DataforgeDatasetConfig = config.dataset
+    dataset: DataforgeDataset = dataset_config.setup()
+    selected: list[SequenceIdentity] = select(dataset.sequences(), config.sequence)
+    print(f"converting {len(selected)} sequence(s)")
+    for identity in selected:
+        target: Path = dataset.convert(identity, force=config.force)
+        if not target.is_file():
+            raise RuntimeError(f"convert produced no recording for {identity.sequence_key}")

packages/dataforge/dataforge/apis/download.py

@@ -0,0 +1,22 @@
+"""``dataforge-download``: fetch (or, for local corpora, verify) a dataset's raw tree."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from dataforge.datasets import AnnotatedDatasetUnion, RobocapConfig
+from dataforge.datasets.base import DataforgeDatasetConfig
+
+
+@dataclass
+class Config:
+    """Download/verify one dataset's raw corpus."""
+
+    dataset: AnnotatedDatasetUnion = field(default_factory=RobocapConfig)
+    """Dataset to download; the raw-tree location lives on the dataset config."""
+
+
+def main(config: Config) -> None:
+    """Run the dataset's download verb."""
+    dataset_config: DataforgeDatasetConfig = config.dataset
+    dataset_config.setup().download()

packages/dataforge/dataforge/apis/register.py

@@ -0,0 +1,52 @@
+"""``dataforge-register``: register base-layer rrds into a local Rerun catalog."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import rerun.blueprint as rrb
+from rerun.catalog import CatalogClient, DatasetEntry, OnDuplicateSegmentLayer
+
+from dataforge import paths
+from dataforge.datasets import AnnotatedDatasetUnion, RobocapConfig, dataset_key
+from dataforge.datasets.base import DataforgeDatasetConfig
+
+LAYER: str = "base"
+"""Only layer dataforge v1 emits."""
+
+
+@dataclass
+class Config:
+    """Register a dataset's converted recordings into a catalog."""
+
+    dataset: AnnotatedDatasetUnion = field(default_factory=RobocapConfig)
+    """Dataset whose rrds get registered; its registry key is the catalog dataset name."""
+    catalog_url: str = "rerun+http://127.0.0.1:51235"
+    """gRPC URL of a locally running ``rerun server`` catalog."""
+
+
+def main(config: Config) -> None:
+    """Create the dataset if needed and register every base-layer rrd (idempotent)."""
+    dataset_config: DataforgeDatasetConfig = config.dataset
+    name: str = dataset_key(dataset_config)
+    rrd_paths: list[Path] = sorted((paths.output_root() / LAYER).glob(f"{name}__*.rrd"))
+    if not rrd_paths:
+        raise FileNotFoundError(f"no {LAYER}-layer rrds for {name} under {paths.output_root() / LAYER}")
+
+    client: CatalogClient = CatalogClient(config.catalog_url)
+    dataset: DatasetEntry = client.create_dataset(name, exist_ok=True)
+    dataset.register(
+        [path.resolve().as_uri() for path in rrd_paths],
+        layer_name=LAYER,
+        on_duplicate=OnDuplicateSegmentLayer.SKIP,
+    ).wait()
+
+    blueprint: rrb.Blueprint | None = dataset_config.setup().default_blueprint()
+    if blueprint is not None:
+        blueprint_path: Path = paths.output_root() / "blueprints" / f"{name}.rbl"
+        blueprint_path.parent.mkdir(parents=True, exist_ok=True)
+        blueprint.save(name, str(blueprint_path))
+        blueprint_path.chmod(0o644)  # uid-squashing NFS mounts can land fresh writes as mode 0000
+        dataset.register_blueprint(blueprint_path.resolve().as_uri(), set_default=True)
+    print(f"registered {len(rrd_paths)} {LAYER}-layer rrds into '{name}' at {config.catalog_url}")

packages/dataforge/dataforge/apis/view.py

@@ -0,0 +1,43 @@
+"""``dataforge-view``: open one converted recording (and its embedded blueprint)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import rerun as rr
+from simplecv.rerun_log_utils import RerunTyroConfig
+
+from dataforge import paths
+from dataforge.datasets import AnnotatedDatasetUnion, RobocapConfig, dataset_key
+from dataforge.datasets.base import DataforgeDatasetConfig
+
+LAYER: str = "base"
+"""Layer whose recordings the viewer loads."""
+
+
+@dataclass
+class Config:
+    """View one converted base-layer recording."""
+
+    rr_config: RerunTyroConfig = field(default_factory=RerunTyroConfig)
+    """Viewer/save/connect behaviour; pass ``--rr-config.headless`` without a display."""
+    dataset: AnnotatedDatasetUnion = field(default_factory=RobocapConfig)
+    """Dataset to view; its registry key prefixes the recording filenames."""
+    sequence: str | None = None
+    """Recording id or sequence key to open; defaults to the first one found."""
+
+
+def main(config: Config) -> None:
+    """Load the selected rrd into the recording stream configured by ``rr_config``."""
+    dataset_config: DataforgeDatasetConfig = config.dataset
+    name: str = dataset_key(dataset_config)
+    candidates: list[Path] = sorted((paths.output_root() / LAYER).glob(f"{name}__*.rrd"))
+    if config.sequence is not None:
+        wanted: str = config.sequence.replace("/", "__")
+        candidates = [path for path in candidates if wanted in path.stem]
+    if not candidates:
+        raise FileNotFoundError(f"no {LAYER}-layer rrd for {name} (sequence={config.sequence}) under {paths.output_root() / LAYER}")
+    rrd_path: Path = candidates[0]
+    print(f"viewing {rrd_path}")
+    rr.log_file_from_path(rrd_path)

packages/dataforge/dataforge/datasets/__init__.py

@@ -0,0 +1,42 @@
+"""Dataset registry: the keys become tyro subcommands on every dataforge verb."""
+
+from __future__ import annotations
+
+import tyro
+
+from dataforge.datasets.base import DataforgeDatasetConfig
+from dataforge.datasets.robocap import RobocapConfig
+
+dataset_defaults: dict[str, DataforgeDatasetConfig] = {
+    "robocap": RobocapConfig(),
+}
+"""Every dataset dataforge knows about, keyed by its CLI name."""
+
+
+def _build_dataset_union() -> type[DataforgeDatasetConfig]:
+    """Build the tyro subcommand union, tolerating a single-entry registry.
+
+    ``subcommand_type_from_defaults`` asserts ``len(defaults) >= 2`` in tyro
+    0.9.x (same workaround as ``robocap_slam.configs.track_dataset_configs``),
+    so with one dataset we hand tyro the concrete config type instead — no
+    placeholder subcommand shows up in ``--help``.
+    """
+    if len(dataset_defaults) >= 2:
+        return tyro.extras.subcommand_type_from_defaults(dataset_defaults, prefix_names=False)
+    only_default: DataforgeDatasetConfig = next(iter(dataset_defaults.values()))
+    return type(only_default)
+
+
+DatasetUnion: type[DataforgeDatasetConfig] = _build_dataset_union()
+"""Config type every verb's ``Config.dataset`` field uses."""
+
+AnnotatedDatasetUnion = tyro.conf.OmitSubcommandPrefixes[DatasetUnion]
+"""``DatasetUnion`` without tyro's per-subcommand flag prefixes."""
+
+
+def dataset_key(config: DataforgeDatasetConfig) -> str:
+    """Registry key of a parsed dataset config (the catalog dataset name)."""
+    for key, default in dataset_defaults.items():
+        if isinstance(config, type(default)):
+            return key
+    raise ValueError(f"Config {type(config).__name__} is not in dataset_defaults")

packages/dataforge/dataforge/datasets/base.py

@@ -0,0 +1,64 @@
+"""Dataset plugin surface: one config dataclass + one dataset class per dataset.
+
+The config/``setup()`` split is simplecv's ``InstantiateConfig`` idiom (see
+``simplecv/configs/base_config.py``): tyro parses the config, ``setup()``
+instantiates the dataset that owns the ``download`` / ``sequences`` /
+``convert`` verbs. dataforge keeps its own copy of the idiom so the dataset
+surface does not inherit simplecv's sequence-loader machinery.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Generic, TypeVar
+
+import rerun.blueprint as rrb
+
+from dataforge.identity import SequenceIdentity
+
+
+@dataclass
+class DataforgeDatasetConfig:
+    """Base config for a dataforge dataset; ``setup()`` builds the dataset."""
+
+    _target: type
+    """Dataset class instantiated by ``setup()``."""
+
+    def setup(self) -> DataforgeDataset[DataforgeDatasetConfig]:
+        """Instantiate the dataset this config describes."""
+        dataset: DataforgeDataset[DataforgeDatasetConfig] = self._target(self)
+        return dataset
+
+
+ConfigT = TypeVar("ConfigT", bound=DataforgeDatasetConfig)
+"""Config type a concrete dataset is parameterized by (simplecv's sequence-loader idiom)."""
+
+
+class DataforgeDataset(Generic[ConfigT], ABC):
+    """A dataset dataforge can download, enumerate, and convert to base-layer rrds."""
+
+    def __init__(self, config: ConfigT) -> None:
+        self.config: ConfigT = config
+
+    @abstractmethod
+    def download(self) -> None:
+        """Fetch (or verify) the raw corpus this dataset converts from."""
+
+    @abstractmethod
+    def sequences(self) -> list[SequenceIdentity]:
+        """Enumerate every convertible sequence found on disk."""
+
+    @abstractmethod
+    def convert(self, identity: SequenceIdentity, *, force: bool) -> Path:
+        """Write the base-layer rrd for one sequence and return its path."""
+
+    def default_blueprint(self) -> rrb.Blueprint | None:
+        """Dataset-wide default blueprint, registered on the catalog dataset.
+
+        Per-recording blueprints are embedded at convert; this is the "dataset
+        default at register" half of the design. ``None`` means the dataset has
+        no sensible corpus-wide layout.
+        """
+        return None