scenezoo 0.1.0__py3-none-any.whl

scenezoo/__init__.py

@@ -0,0 +1,43 @@
+"""A small, extensible toolkit for 3D dataset preprocessing."""
+
+from .dataset import (
+    CheckIssue,
+    DataFormatError,
+    Dataset,
+    DatasetCheckError,
+    DatasetCheckReport,
+    DatasetInfo,
+    DatasetSpec,
+    FrameBatch,
+    FrameSourceSpec,
+    PointBatch,
+    SceneZooError,
+    Segmentation3D,
+    UnsupportedOperationError,
+    get_dataset,
+    get_dataset_info,
+    list_datasets,
+    register_dataset,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "CheckIssue",
+    "DataFormatError",
+    "Dataset",
+    "DatasetCheckError",
+    "DatasetCheckReport",
+    "DatasetInfo",
+    "DatasetSpec",
+    "FrameBatch",
+    "FrameSourceSpec",
+    "PointBatch",
+    "SceneZooError",
+    "Segmentation3D",
+    "UnsupportedOperationError",
+    "get_dataset",
+    "get_dataset_info",
+    "list_datasets",
+    "register_dataset",
+]

scenezoo/config.py

@@ -0,0 +1,6 @@
+import os
+
+XDG_CACHE_HOME = os.getenv("XDG_CACHE_HOME", "~/.cache")
+SCENEZOO_CACHE_DIR = os.path.expanduser(
+    os.getenv("SCENEZOO_CACHE_DIR", os.path.join(XDG_CACHE_HOME, "scenezoo"))
+)

scenezoo/dataset/__init__.py

@@ -0,0 +1,52 @@
+"""Dataset loading, registration, and shared types."""
+
+from .base import (
+    DataFormatError,
+    Dataset,
+    DatasetCheckError,
+    SceneZooError,
+    UnsupportedOperationError,
+)
+from .check import CheckIssue, DatasetCheckReport
+from .registry import (
+    DatasetInfo,
+    DatasetRegistry,
+    _REGISTRY,
+    get_dataset,
+    get_dataset_info,
+    is_registered,
+    list_datasets,
+    register_dataset,
+)
+from .spec import DatasetSpec, FrameSourceSpec
+from .types import FrameBatch, PointBatch, Segmentation3D
+
+from .builtins import register_builtin_datasets
+
+# This package exposes adapter classes lazily through ``__getattr__``.
+from . import scene as scene
+
+register_builtin_datasets(_REGISTRY)
+
+__all__ = [
+    "DataFormatError",
+    "CheckIssue",
+    "Dataset",
+    "DatasetCheckError",
+    "DatasetCheckReport",
+    "DatasetInfo",
+    "DatasetRegistry",
+    "DatasetSpec",
+    "FrameBatch",
+    "FrameSourceSpec",
+    "PointBatch",
+    "SceneZooError",
+    "Segmentation3D",
+    "UnsupportedOperationError",
+    "get_dataset",
+    "get_dataset_info",
+    "is_registered",
+    "list_datasets",
+    "register_dataset",
+    "scene",
+]

scenezoo/dataset/base.py

@@ -0,0 +1,247 @@
+"""Core dataset abstraction."""
+
+from __future__ import annotations
+
+from functools import cached_property
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, ClassVar, Sequence
+
+from .spec import DatasetSpec
+
+if TYPE_CHECKING:
+    from .check import DatasetCheckReport
+
+
+class SceneZooError(Exception):
+    """Base exception raised by scenezoo."""
+
+
+class UnsupportedOperationError(SceneZooError):
+    """Raised when a dataset does not provide a requested operation or item."""
+
+
+class DataFormatError(SceneZooError):
+    """Raised when source data does not match its documented format."""
+
+
+class DatasetCheckError(SceneZooError):
+    """Raised when an explicitly requested dataset check finds errors."""
+
+    def __init__(self, report: DatasetCheckReport) -> None:
+        self.report = report
+        super().__init__(report.format())
+
+
+class Dataset:
+    """Minimal base class for 3D scene datasets.
+
+    Optional operations are regular methods rather than abstract methods. This
+    lets partial adapters expose exactly the data they really provide.
+    """
+
+    name: ClassVar[str] = "dataset"
+    spec: ClassVar[DatasetSpec] = DatasetSpec()
+    _OPERATIONS: ClassVar[dict[str, str]] = {
+        "mesh": "get_mesh",
+        "points": "get_points",
+        "segmentation": "get_segmentation",
+        "boxes": "get_boxes",
+        "frames": "get_frames",
+    }
+
+    def __init__(
+        self,
+        root_dir: str | Path,
+        *,
+        cache_dir: str | Path | None = None,
+        offline: bool = False,
+        invalid_obj_id: int = -1,
+    ) -> None:
+        self.root_dir = Path(root_dir).expanduser()
+        self.cache_dir = None if cache_dir is None else Path(cache_dir).expanduser()
+        self.offline = offline
+        self.invalid_obj_id = invalid_obj_id
+
+    @cached_property
+    def splits(self) -> dict[str, list[str]]:
+        """Sample IDs of every named split, loaded once on first access."""
+
+        splits = self._load_splits()
+        if not isinstance(splits, dict):
+            raise DataFormatError("_load_splits() must return a dictionary.")
+        return splits
+
+    @cached_property
+    def metadata(self) -> dict[str, Any]:
+        """Dataset-wide metadata such as label maps, loaded once on first access."""
+
+        metadata = self._load_metadata()
+        if not isinstance(metadata, dict):
+            raise DataFormatError("_load_metadata() must return a dictionary.")
+        return metadata
+
+    def _load_splits(self) -> dict[str, list[str]]:
+        raise NotImplementedError
+
+    def _load_metadata(self) -> dict[str, Any]:
+        return {}
+
+    def get_ids(self, split: str) -> list[str]:
+        """Return the sample IDs of ``split``; see :attr:`splits` for the names."""
+
+        try:
+            return self.splits[split]
+        except KeyError as exc:
+            available = ", ".join(sorted(self.splits))
+            raise KeyError(
+                f"Unknown split {split!r}. Available splits: {available}"
+            ) from exc
+
+    def open_file(self, file_path: str | Path, mode: str = "rb"):
+        """Open a local or cached remote metadata file using this dataset's policy."""
+        from ..io.cache import open_cached_file
+
+        return open_cached_file(
+            file_path,
+            mode,
+            cache_dir=self.cache_dir,
+            offline=self.offline,
+        )
+
+    def check(
+        self,
+        *,
+        sample_ids: Sequence[str] | None = None,
+        require_complete: bool = False,
+        raise_on_error: bool = False,
+    ) -> DatasetCheckReport:
+        """Explicitly inspect the dataset path and return an actionable report.
+
+        Checks are never run by the constructor or by data access methods. Use
+        ``sample_ids`` for a focused check, or ``require_complete=True`` to treat
+        missing scenes from declared splits as errors rather than warnings.
+        """
+
+        report = self._check(
+            sample_ids=sample_ids,
+            require_complete=require_complete,
+        )
+        if raise_on_error:
+            report.raise_for_errors()
+        return report
+
+    def _check(
+        self,
+        *,
+        sample_ids: Sequence[str] | None,
+        require_complete: bool,
+    ) -> DatasetCheckReport:
+        from .check import CheckBuilder
+
+        builder = CheckBuilder(self.name, self.root_dir)
+        builder.require_directory(
+            self.root_dir,
+            code="missing-root",
+            message="The configured dataset root does not exist or is not a directory.",
+            expected="An existing directory containing the dataset files.",
+            hint="Correct root_dir or download and extract the dataset first.",
+        )
+        return builder.finish()
+
+    @classmethod
+    def supports(cls, operation: str) -> bool:
+        """Whether this adapter implements a standard operation.
+
+        ``operation`` is one of ``mesh``, ``points``, ``segmentation``,
+        ``boxes``, or ``frames``.
+        """
+
+        try:
+            method_name = cls._OPERATIONS[operation]
+        except KeyError as exc:
+            valid = ", ".join(sorted(cls._OPERATIONS))
+            raise ValueError(
+                f"Unknown operation {operation!r}. Valid operations: {valid}"
+            ) from exc
+        return getattr(cls, method_name) is not getattr(Dataset, method_name)
+
+    @classmethod
+    def capabilities(cls) -> tuple[str, ...]:
+        """The standard operations this adapter implements."""
+
+        return tuple(name for name in cls._OPERATIONS if cls.supports(name))
+
+    def _unsupported(self, operation: str) -> None:
+        raise UnsupportedOperationError(
+            f"{self.__class__.__name__} does not support the {operation!r} operation."
+        )
+
+    def get_mesh(self, sample_id: str, *, mesh_type: str | None = None):
+        """Return an Open3D ``TriangleMesh``.
+
+        ``mesh_type`` selects a dataset-specific variant listed in
+        ``spec.mesh_types``; ``None`` selects the adapter's default.
+        """
+
+        self._unsupported("mesh")
+
+    def get_points(self, sample_id: str):
+        """Return a :class:`~scenezoo.PointBatch` with row-aligned attributes."""
+
+        self._unsupported("points")
+
+    def get_segmentation(self, sample_id: str):
+        """Return :class:`~scenezoo.Segmentation3D` labels.
+
+        Labels are defined on the ``domain`` the result reports (points, mesh
+        vertices, or mesh faces) and align with the matching geometry.
+        """
+
+        self._unsupported("segmentation")
+
+    def get_boxes(self, sample_id: str, *, box_type: str = "mobb_gravity"):
+        """Return ``(boxes_by_id, id_to_name)`` with Open3D bounding boxes.
+
+        ``box_type`` is ``obb_gt`` for released annotations, or ``aabb``,
+        ``obb``, ``mobb``, or ``mobb_gravity`` to fit boxes to annotated points.
+        """
+
+        self._unsupported("boxes")
+
+    def get_frames(
+        self,
+        sample_id: str,
+        *,
+        indices=None,
+        step: int = 1,
+        items=("rgb", "rgb_intrinsics", "world_to_camera"),
+        output_size=None,
+        center_crop: bool = False,
+        rotate_to_up: bool = True,
+    ):
+        """Return synchronized frames as a :class:`~scenezoo.FrameBatch`.
+
+        Args:
+            indices: Source frame numbers to read, in the requested order.
+                ``None`` selects every frame with a valid camera.
+            step: Sampling stride when ``indices`` is ``None``.
+            items: :class:`~scenezoo.FrameBatch` fields to load.
+            output_size: Optional ``(width, height)`` for every image item.
+            center_crop: Preserve the aspect ratio by resizing then cropping.
+            rotate_to_up: Rotate portrait captures upright; the camera roll is
+                applied to ``world_to_camera`` so intrinsics stay pinhole.
+        """
+
+        self._unsupported("frames")
+
+    def __repr__(self) -> str:
+        loaded_splits = self.__dict__.get("splits")
+        split_sizes = (
+            None
+            if loaded_splits is None
+            else {key: len(value) for key, value in loaded_splits.items()}
+        )
+        return (
+            f"{self.__class__.__name__}(root_dir={str(self.root_dir)!r}, "
+            f"splits={split_sizes})"
+        )

scenezoo/dataset/builtins.py

@@ -0,0 +1,320 @@
+"""Import-light catalog for built-in dataset adapters."""
+
+from __future__ import annotations
+
+from .spec import DatasetSpec, FrameSourceSpec
+
+_RGBD_ITEMS = (
+    "rgb",
+    "depth",
+    "rgb_intrinsics",
+    "depth_intrinsics",
+    "world_to_camera",
+)
+
+
+BUILTIN_DATASETS = (
+    {
+        "name": "scannetv2",
+        "aliases": ("scannet",),
+        "import_path": "scenezoo.dataset.scene.scannet:ScanNet",
+        "description": "ScanNet v2 indoor RGB-D scenes.",
+        "spec": DatasetSpec(
+            sample_unit="scene",
+            operations=("mesh", "segmentation", "boxes", "frames"),
+            mesh_types=("raw", "decimated", "semantic"),
+            segmentation_types=("instance", "semantic"),
+            box_types=("aabb", "obb", "mobb", "mobb_gravity"),
+            label_spaces=("scannet", "nyu40"),
+            frame_sources=(
+                FrameSourceSpec(
+                    "sens",
+                    _RGBD_ITEMS
+                    + ("instance_maps", "semantic_maps", "timestamps", "imu"),
+                    ("PINHOLE",),
+                ),
+            ),
+        ),
+    },
+    {
+        "name": "scannetppv2",
+        "aliases": ("scannet++", "scannetpp", "scannet++v2"),
+        "import_path": "scenezoo.dataset.scene.scannetpp:ScanNetPP",
+        "description": "ScanNet++ v2 scans with iPhone, DSLR, and panoramic captures.",
+        "spec": DatasetSpec(
+            sample_unit="scene",
+            operations=("mesh", "points", "segmentation", "boxes", "frames"),
+            mesh_types=("raw", "semantic"),
+            segmentation_types=("instance", "semantic"),
+            box_types=("aabb", "obb", "mobb", "mobb_gravity"),
+            frame_sources=(
+                FrameSourceSpec(
+                    "iphone",
+                    _RGBD_ITEMS + ("frame_mask", "timestamps", "imu"),
+                    ("PINHOLE",),
+                ),
+                FrameSourceSpec(
+                    "iphone_colmap",
+                    _RGBD_ITEMS + ("frame_mask", "timestamps", "imu"),
+                    ("SIMPLE_PINHOLE", "PINHOLE", "OPENCV", "OPENCV_FISHEYE"),
+                ),
+                FrameSourceSpec(
+                    "dslr",
+                    (
+                        "rgb",
+                        "rgb_intrinsics",
+                        "world_to_camera",
+                        "frame_mask",
+                        "is_bad",
+                    ),
+                    ("SIMPLE_PINHOLE", "PINHOLE", "OPENCV", "OPENCV_FISHEYE"),
+                ),
+                FrameSourceSpec(
+                    "dslr_undistorted",
+                    (
+                        "rgb",
+                        "rgb_intrinsics",
+                        "world_to_camera",
+                        "frame_mask",
+                        "is_bad",
+                    ),
+                    ("PINHOLE",),
+                ),
+                FrameSourceSpec(
+                    "dslr_original",
+                    (
+                        "rgb",
+                        "rgb_intrinsics",
+                        "world_to_camera",
+                        "frame_mask",
+                        "is_bad",
+                    ),
+                    ("SIMPLE_PINHOLE", "PINHOLE", "OPENCV", "OPENCV_FISHEYE"),
+                ),
+                FrameSourceSpec(
+                    "panocam",
+                    (
+                        "rgb",
+                        "depth",
+                        "world_to_camera",
+                        "frame_mask",
+                        "azimuth",
+                        "elevation",
+                    ),
+                    ("EQUIRECTANGULAR",),
+                ),
+                FrameSourceSpec(
+                    "panocam_resized",
+                    (
+                        "rgb",
+                        "depth",
+                        "world_to_camera",
+                        "frame_mask",
+                        "azimuth",
+                        "elevation",
+                    ),
+                    ("EQUIRECTANGULAR",),
+                ),
+            ),
+        ),
+    },
+    {
+        "name": "arkitscenes",
+        "aliases": (),
+        "import_path": "scenezoo.dataset.scene.arkitscenes:ARKitScenes",
+        "description": "ARKitScenes scans, mobile RGB-D streams, and FARO captures.",
+        "spec": DatasetSpec(
+            sample_unit="scene",
+            operations=("mesh", "segmentation", "boxes", "frames"),
+            mesh_types=("raw",),
+            segmentation_types=("instance",),
+            box_types=("obb_gt",),
+            frame_sources=tuple(
+                FrameSourceSpec(
+                    name,
+                    _RGBD_ITEMS + ("confidence_maps", "timestamps"),
+                    ("PINHOLE",),
+                )
+                for name in ("mov", "lowres_wide", "wide")
+            )
+            + tuple(
+                FrameSourceSpec(
+                    name,
+                    ("rgb", "rgb_intrinsics", "world_to_camera", "timestamps"),
+                    ("PINHOLE",),
+                )
+                for name in ("ultrawide", "vga_wide")
+            )
+            + (
+                FrameSourceSpec(
+                    "threedod",
+                    _RGBD_ITEMS + ("timestamps",),
+                    ("PINHOLE",),
+                ),
+            ),
+        ),
+    },
+    {
+        "name": "multiscan",
+        "aliases": (),
+        "import_path": "scenezoo.dataset.scene.multiscan:MultiScan",
+        "description": "MultiScan indoor scenes.",
+        "spec": DatasetSpec(
+            sample_unit="scene",
+            operations=("mesh", "segmentation", "boxes", "frames"),
+            mesh_types=("ply", "textured"),
+            segmentation_types=(
+                "object_instance",
+                "part_instance",
+                "object_semantic",
+                "part_semantic",
+            ),
+            box_types=("aabb", "obb", "mobb", "mobb_gravity", "obb_gt"),
+            frame_sources=(
+                FrameSourceSpec(
+                    "capture",
+                    _RGBD_ITEMS
+                    + ("confidence_maps", "timestamps", "exposure_durations"),
+                    ("PINHOLE",),
+                ),
+            ),
+        ),
+    },
+    {
+        "name": "ase",
+        "aliases": ("aria-ase", "aria-synthetic-environments"),
+        "import_path": ("scenezoo.dataset.scene.ase:AriaSyntheticEnvironments"),
+        "description": "Aria Synthetic Environments egocentric synthetic scenes.",
+        "spec": DatasetSpec(
+            sample_unit="sequence",
+            operations=("points", "boxes", "frames"),
+            box_types=("obb_gt",),
+            frame_sources=(
+                FrameSourceSpec(
+                    "rgb",
+                    _RGBD_ITEMS + ("instance_maps", "frame_mask", "timestamps"),
+                    ("FISHEYE624",),
+                ),
+            ),
+        ),
+    },
+    {
+        "name": "3rscan",
+        "aliases": ("threerscan",),
+        "import_path": "scenezoo.dataset.scene.threerscan:ThreeRScan",
+        "description": "3RScan indoor scene sequences.",
+        "spec": DatasetSpec(
+            sample_unit="scene",
+            operations=("mesh", "segmentation", "boxes", "frames"),
+            mesh_types=("textured", "instance", "instance_aligned"),
+            segmentation_types=(
+                "instance",
+                "global",
+                "nyu40",
+                "eigen13",
+                "rio27",
+                "rio7",
+            ),
+            box_types=("aabb", "obb", "mobb", "mobb_gravity", "obb_gt"),
+            label_spaces=("global", "nyu40", "eigen13", "rio27", "rio7"),
+            frame_sources=(FrameSourceSpec("sequence", _RGBD_ITEMS, ("PINHOLE",)),),
+        ),
+    },
+    {
+        "name": "s3dis",
+        "aliases": ("stanford3d", "stanford-3d"),
+        "import_path": "scenezoo.dataset.scene.s3dis:S3DIS",
+        "description": "S3DIS room point clouds, point-domain labels, and boxes.",
+        "spec": DatasetSpec(
+            sample_unit="room",
+            operations=("points", "segmentation", "boxes"),
+            segmentation_types=("instance", "semantic"),
+            box_types=("aabb", "obb", "mobb", "mobb_gravity"),
+            label_spaces=("s3dis13",),
+        ),
+    },
+    {
+        "name": "matterport3d",
+        "aliases": (),
+        "import_path": "scenezoo.dataset.scene.matterport3d:Matterport3D",
+        "description": "Matterport3D houses with region annotations and RGB-D views.",
+        "spec": DatasetSpec(
+            sample_unit="scene",
+            operations=("mesh", "segmentation", "boxes", "frames"),
+            mesh_types=("raw", "house", "region", "poisson"),
+            segmentation_types=("instance", "semantic"),
+            box_types=("obb_gt", "aabb", "obb", "mobb", "mobb_gravity"),
+            label_spaces=("raw", "mpcat40", "nyu40", "eigen13"),
+            frame_sources=(
+                FrameSourceSpec(
+                    "raw",
+                    _RGBD_ITEMS,
+                    ("OPENCV",),
+                ),
+                FrameSourceSpec(
+                    "undistorted",
+                    _RGBD_ITEMS,
+                    ("PINHOLE",),
+                ),
+            ),
+        ),
+    },
+    {
+        "name": "scenenn",
+        "aliases": (),
+        "import_path": "scenezoo.dataset.scene.scenenn:SceneNN",
+        "description": "SceneNN meshes and vertex-domain instance segmentation.",
+        "spec": DatasetSpec(
+            sample_unit="scene",
+            operations=("mesh", "segmentation"),
+            mesh_types=("raw", "instance"),
+            segmentation_types=("instance",),
+        ),
+    },
+    {
+        "name": "structured3d",
+        "aliases": ("structure3d",),
+        "import_path": "scenezoo.dataset.scene.structured3d:Structured3D",
+        "description": "Structured3D renderings, layouts, boxes, and fused points.",
+        "spec": DatasetSpec(
+            sample_unit="scene",
+            operations=("mesh", "points", "segmentation", "boxes", "frames"),
+            mesh_types=("layout",),
+            segmentation_types=("instance", "semantic"),
+            box_types=("obb_gt",),
+            label_spaces=("nyu40", "pointcept25"),
+            frame_sources=(
+                FrameSourceSpec(
+                    "perspective",
+                    _RGBD_ITEMS
+                    + ("semantic_maps", "instance_maps", "albedo", "normal_maps"),
+                    ("PINHOLE",),
+                    ("empty", "full"),
+                    ("room_id",),
+                ),
+                FrameSourceSpec(
+                    "panorama",
+                    (
+                        "rgb",
+                        "depth",
+                        "world_to_camera",
+                        "semantic_maps",
+                        "instance_maps",
+                        "albedo",
+                        "normal_maps",
+                    ),
+                    ("EQUIRECTANGULAR",),
+                    ("empty", "simple", "full"),
+                    ("room_id",),
+                ),
+            ),
+        ),
+    },
+)
+
+
+def register_builtin_datasets(registry) -> None:
+    """Populate a registry without importing any adapter modules."""
+
+    for entry in BUILTIN_DATASETS:
+        registry.register_lazy(**entry)