diff --git a/docs/CONVENTIONS.md b/docs/CONVENTIONS.md index 7506ac3df..bbeea9edb 100644 --- a/docs/CONVENTIONS.md +++ b/docs/CONVENTIONS.md @@ -22,6 +22,39 @@ EVA episode extrinsics store one `base_T_cam` matrix for each arm. Human episodes store the head pose as `world_T_head`. The head frame and the camera frame are the same frame for egocentric human data. +### The calibration block + +An episode written after the `calibration` attribute existed names one +reference frame and expresses every pose in it. + +```text +calibration.reference_frame robot_base | slam_world | camera: +calibration.cameras[c].ref_T_cam camera c's pose in the reference frame +calibration.arm_bases[side] ref_T_armbase, the arm base pose in it +``` + +The camera that defines the reference frame needs no `ref_T_cam`. It is the +identity by definition. + +Each camera also declares its projection model and distortion coefficients. + +```text +calibration.cameras[c].model PINHOLE | OPENCV | KANNALA_BRANDT +calibration.cameras[c].distortion coefficients in that model's order +calibration.cameras[c].rectified whether the stored frames are rectified +``` + +A camera that declares no model is `PINHOLE` with no coefficients. No +projection site honors a non-pinhole model yet. Collect the declaration anyway: +it measures a vendor's rig, and a rig that has moved cannot be recalibrated +after the fact. + +`Calibration.base_T_cam(side)` composes the two and returns what the EVA +transform pipeline consumes. An episode that predates the block reaches the +same value through the shim in `egomimic/rldb/zarr/calibration.py`: its +reference frame is `camera:front_1`, so `arm_bases[side]` is the inverse of +the stored `extrinsics[side]`. + ### Existing episodes need no migration The `ref_T_cam` naming records the direction the code already used. It changes diff --git a/egomimic/rldb/embodiment/embodiment.py b/egomimic/rldb/embodiment/embodiment.py index b8a802028..b4c0ec3bd 100644 --- a/egomimic/rldb/embodiment/embodiment.py +++ b/egomimic/rldb/embodiment/embodiment.py @@ -135,6 +135,36 @@ def action_space(self) -> str: ) return spaces.pop() + @property + def arity(self) -> str | None: + """Return the arity suffix encoded in ``embodiment_name``. + + Returns: + The text after ``_``, or ``None`` when + the resolved value has no matching embodiment name. + """ + if self.embodiment_name is None: + return None + prefix = f"{self.platform.embodiment_prefix}_" + if self.embodiment_name.startswith(prefix): + return self.embodiment_name[len(prefix) :] + return None + + @property + def sides(self) -> tuple[str, ...]: + """Return side candidates implied by the encoded arity. + + ``left_arm`` and ``right_arm`` select one side. All other values return + both candidates; callers resolving a morphology mapping must still + filter candidates not present in ``end_effectors``. + """ + arity = self.arity + if arity == "left_arm": + return ("left",) + if arity == "right_arm": + return ("right",) + return SIDES + @property def embodiment_class(self) -> type["Embodiment"]: """Import and return the platform's configured ``Embodiment`` class. diff --git a/egomimic/rldb/embodiment/eva.py b/egomimic/rldb/embodiment/eva.py index d63df477f..0f9b07498 100644 --- a/egomimic/rldb/embodiment/eva.py +++ b/egomimic/rldb/embodiment/eva.py @@ -26,7 +26,15 @@ class Eva(Embodiment): + """Dataset transforms and visualization for the EVA X5 platform. + + ``EXTRINSICS`` is used only when a sample has no episode-specific + ``base_T_cam`` poses. ``INTRINSICS`` is the visualization fallback when a + sample has no camera matrix. Both constants describe the original EVA rig. + """ + INTRINSICS = ARIA_INTRINSICS + #: Compatibility fallback: one ``base_T_cam`` pose per arm. EXTRINSICS = { "left": np.array( [ @@ -248,6 +256,7 @@ def _build_eva_bimanual_eef_frame_transform_list( ) -> list[Transform]: """EVA bimanual transform pipeline with actions expressed relative to the current EEF pose (wrist frame), analogous to keypoints relative to wrist pose.""" + # Supply class-default poses only for sample keys absent at transform time. extrinsics = Eva.EXTRINSICS left_base_T_cam_pose = _matrix_to_xyzwxyz(extrinsics["left"][None, :])[0] right_base_T_cam_pose = _matrix_to_xyzwxyz(extrinsics["right"][None, :])[0] @@ -411,6 +420,7 @@ def _build_eva_bimanual_transform_list( is_quat: bool = True, ) -> list[Transform]: """Canonical EVA bimanual transform pipeline used by tests and notebooks.""" + # Supply class-default poses only for sample keys absent at transform time. extrinsics = Eva.EXTRINSICS left_base_T_cam_pose = _matrix_to_xyzwxyz(extrinsics["left"][None, :])[0] right_base_T_cam_pose = _matrix_to_xyzwxyz(extrinsics["right"][None, :])[0] diff --git a/egomimic/rldb/zarr/action_chunk_transforms.py b/egomimic/rldb/zarr/action_chunk_transforms.py index 9daf5e057..6c5292ea1 100644 --- a/egomimic/rldb/zarr/action_chunk_transforms.py +++ b/egomimic/rldb/zarr/action_chunk_transforms.py @@ -35,6 +35,18 @@ xyzw_to_wxyz, ) + +def base_T_cam_pose_key(side: str) -> str: + """Return the sample key for ``base_T_cam`` as ``[xyz, qw, qx, qy, qz]``.""" + return f"{side}_base_T_cam_pose" + + +def _apply_fallbacks(batch, extra_batch_key) -> None: + """Add fallback values for absent keys without replacing sample values.""" + for key, value in (extra_batch_key or {}).items(): + batch.setdefault(key, value) + + # --------------------------------------------------------------------------- # Base Transform # --------------------------------------------------------------------------- @@ -159,7 +171,9 @@ def __init__( frame. This pose represents ``reference_T_target``. chunk_world: Batch key for the input poses. transformed_key_name: Batch key for the output poses. - extra_batch_key: Values to add to the batch before the transform. + extra_batch_key: Fallback values for keys the batch does not + already carry. A per-episode value in the batch wins, so a + class constant here is a default and not an override. mode: Input and output layout. Use ``"xyz"``, ``"xyzypr"``, or ``"xyzwxyz"``. inverse: If true, compute ``target_T_chunk`` from @@ -190,7 +204,7 @@ def transform(self, batch): ValueError: If ``mode`` is not a supported pose layout. """ # Flatten the leading chunk dimensions into one pose dimension. - batch.update(self.extra_batch_key or {}) + _apply_fallbacks(batch, self.extra_batch_key) target_world_pose = np.asarray(batch[self.target_world]) chunk_world_poses = np.asarray(batch[self.chunk_world]) chunk_world_poses_shape = None @@ -437,7 +451,7 @@ def transform(self, batch): returns batch with new key containing transformed chunk world in target frame: (T, 14) """ - batch.update(self.extra_batch_key or {}) + _apply_fallbacks(batch, self.extra_batch_key) left_target_world = batch[self.left_target_world] right_target_world = batch[self.right_target_world] chunk_world = batch[self.chunk_world] diff --git a/egomimic/rldb/zarr/calibration.py b/egomimic/rldb/zarr/calibration.py new file mode 100644 index 000000000..cae947c14 --- /dev/null +++ b/egomimic/rldb/zarr/calibration.py @@ -0,0 +1,578 @@ +"""Parse and serialize per-episode camera and arm-base calibration. + +Calibration is a physical measurement of the rig that recorded one episode, so +it travels with the episode instead of living in a class constant. An episode +stores it under one ``calibration`` attribute:: + + "calibration": { + "reference_frame": "camera:front_1", + "cameras": { + "front_1": { + "K": [[fx, 0, cx, 0], [0, fy, cy, 0], [0, 0, 1, 0]], + "model": "PINHOLE", + "distortion": [], + "resolution": [W, H], + "rectified": true, + "ref_T_cam": [[...]], + }, + }, + "arm_bases": {"left": [[...]], "right": [[...]]}, + } + +Every rigid transform follows ``docs/CONVENTIONS.md``: ``A_T_B`` maps +coordinates from frame ``B`` to frame ``A``. ``ref_T_cam`` is the camera pose +in the reference frame, and ``arm_bases[side]`` is ``ref_T_armbase``, the arm +base pose in the reference frame. + +Episodes written before this block existed store ``intrinsics`` and +``extrinsics`` at the top level instead. :func:`read_calibration` reads either +form, so no stored episode needs a rewrite. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any + +import numpy as np + +#: Conventional camera name assigned to bare legacy intrinsics and extrinsics. +LEGACY_REFERENCE_CAMERA = "front_1" + +#: Supported reference frames that do not name a camera. +STATIC_REFERENCE_FRAMES = frozenset({"robot_base", "slam_world"}) + +#: Prefix for a camera reference frame, as in ``camera:front_1``. +CAMERA_FRAME_PREFIX = "camera:" + +#: Prefix of an image-array key; the suffix is its camera name. +IMAGE_KEY_PREFIX = "images." + +#: Supported projection models mapped to accepted distortion-vector lengths. +CAMERA_MODELS = { + "PINHOLE": frozenset({0}), + "OPENCV": frozenset({4, 5, 8, 12, 14}), + "KANNALA_BRANDT": frozenset({4}), +} + +#: Projection model used when a camera omits ``model``. +DEFAULT_CAMERA_MODEL = "PINHOLE" + +_CAMERA_FIELDS = frozenset( + {"K", "model", "distortion", "resolution", "rectified", "ref_T_cam"} +) +_CALIBRATION_FIELDS = frozenset({"reference_frame", "cameras", "arm_bases"}) + + +class CalibrationError(ValueError): + """Raised when current or legacy calibration metadata is malformed.""" + + +def _matrix(value: Any, shape: tuple[int, ...], where: str) -> np.ndarray: + """Convert ``value`` to a finite float64 array with exactly ``shape``.""" + try: + arr = np.asarray(value, dtype=np.float64) + except (TypeError, ValueError) as exc: + raise CalibrationError(f"{where}: not a numeric array ({exc})") from exc + if arr.shape != shape: + raise CalibrationError( + f"{where}: expected shape {shape}, got {arr.shape}" + ) + if not np.isfinite(arr).all(): + raise CalibrationError(f"{where}: contains a non-finite value") + return arr + + +def _camera_matrix(value: Any, where: str) -> np.ndarray: + """Return finite ``[K_3x3 | 0]``; append a zero column to a 3×3 input.""" + try: + arr = np.asarray(value, dtype=np.float64) + except (TypeError, ValueError) as exc: + raise CalibrationError(f"{where}: not a numeric array ({exc})") from exc + if arr.shape == (3, 3): + arr = np.hstack([arr, np.zeros((3, 1))]) + if arr.shape != (3, 4): + raise CalibrationError( + f"{where}: expected a 3x4 camera matrix (pad a bare 3x3 with " + f"np.hstack([K, np.zeros((3, 1))])), got shape {arr.shape}" + ) + if not np.isfinite(arr).all(): + raise CalibrationError(f"{where}: contains a non-finite value") + return arr + + +@dataclass(frozen=True) +class CameraCalibration: + """Normalized calibration metadata for one camera stream. + + Attributes: + name: The camera name. It matches the ``images.`` array key. + K: The 3×4 ``[K_3x3 | 0]`` matrix, or ``None`` when undeclared. + model: A key in ``CAMERA_MODELS``. Current projection code does not + branch on this value. + distortion: The model-specific distortion vector. Its length is + validated against ``CAMERA_MODELS``; current projection code does + not consume it. + resolution: ``(width, height)`` in pixels, or ``None``. + rectified: Whether the stored frames are already rectified. + ref_T_cam: The 4×4 camera pose in the episode reference frame, or + ``None`` when the episode does not state it. + """ + + name: str + K: np.ndarray | None = None + model: str = DEFAULT_CAMERA_MODEL + distortion: tuple[float, ...] = () + resolution: tuple[int, int] | None = None + rectified: bool = True + ref_T_cam: np.ndarray | None = None + + def to_jsonable(self) -> dict[str, Any]: + """Serialize this camera to plain values accepted by Zarr attributes.""" + out: dict[str, Any] = { + "model": self.model, + "rectified": bool(self.rectified), + } + if self.K is not None: + out["K"] = self.K.tolist() + if self.distortion: + out["distortion"] = [float(c) for c in self.distortion] + if self.resolution is not None: + out["resolution"] = [int(self.resolution[0]), int(self.resolution[1])] + if self.ref_T_cam is not None: + out["ref_T_cam"] = self.ref_T_cam.tolist() + return out + + +@dataclass(frozen=True) +class Calibration: + """Normalized calibration metadata for one episode. + + Attributes: + reference_frame: ``robot_base``, ``slam_world``, or + ``camera:``. Every rigid pose in this block is + expressed in this frame. + cameras: A mapping from camera name to its calibration. + arm_bases: A mapping from ``"left"`` or ``"right"`` to a 4×4 + ``ref_T_armbase`` matrix. Human episodes leave this empty. + legacy: Whether :func:`read_calibration` built this from the + ``intrinsics`` and ``extrinsics`` attributes of an older episode. + """ + + reference_frame: str + cameras: Mapping[str, CameraCalibration] = field(default_factory=dict) + arm_bases: Mapping[str, np.ndarray] = field(default_factory=dict) + legacy: bool = False + + @property + def reference_camera(self) -> str | None: + """Return the camera that defines the reference frame, if any.""" + if self.reference_frame.startswith(CAMERA_FRAME_PREFIX): + return self.reference_frame[len(CAMERA_FRAME_PREFIX) :] + return None + + def default_camera(self) -> str | None: + """Select a camera when the caller does not name one. + + Selection order is the declared reference camera, ``front_1``, the + first name containing ``front`` (case-insensitive), and finally the + first declared camera. + """ + reference = self.reference_camera + if reference is not None and reference in self.cameras: + return reference + if LEGACY_REFERENCE_CAMERA in self.cameras: + return LEGACY_REFERENCE_CAMERA + front = next((n for n in self.cameras if "front" in n.lower()), None) + if front is not None: + return front + return next(iter(self.cameras), None) + + def K(self, camera: str | None = None) -> np.ndarray | None: + """Return one camera's normalized 3×4 ``[K_3x3 | 0]`` matrix. + + Args: + camera: A camera name. ``None`` selects :meth:`default_camera`. + + Returns: + The camera matrix, or ``None`` if the camera is absent or carries + no ``K``. + """ + name = camera if camera is not None else self.default_camera() + if name is None: + return None + entry = self.cameras.get(name) + return None if entry is None else entry.K + + def ref_T_cam(self, camera: str | None = None) -> np.ndarray | None: + """Return one camera's pose in the episode reference frame. + + A missing pose for the reference camera is synthesized as identity; + another missing pose returns ``None``. + """ + name = camera if camera is not None else self.default_camera() + if name is None: + return None + entry = self.cameras.get(name) + if entry is not None and entry.ref_T_cam is not None: + return entry.ref_T_cam + if name == self.reference_camera: + return np.eye(4) + return None + + def base_T_cam(self, side: str, camera: str | None = None) -> np.ndarray | None: + """Compose one camera's pose in the selected arm-base frame. + + This is the quantity the EVA transform pipeline consumes, and it is + what the ``extrinsics`` attribute of a legacy episode stores directly. + + Args: + side: ``"left"`` or ``"right"``. + camera: A camera name. ``None`` selects :meth:`default_camera`. + + Returns: + A 4×4 ``base_T_cam`` matrix, or ``None`` when the selected arm base + or camera pose is unavailable. + """ + ref_T_armbase = self.arm_bases.get(side) + ref_T_cam = self.ref_T_cam(camera) + if ref_T_armbase is None or ref_T_cam is None: + return None + return np.linalg.inv(ref_T_armbase) @ ref_T_cam + + def intrinsics(self) -> dict[str, np.ndarray]: + """Return ``{camera_name: K}`` for cameras with intrinsic matrices.""" + return { + name: cam.K for name, cam in self.cameras.items() if cam.K is not None + } + + def extrinsics(self, camera: str | None = None) -> dict[str, np.ndarray]: + """Return composable ``{side: base_T_cam}`` values in legacy layout.""" + out = {} + for side in self.arm_bases: + base_T_cam = self.base_T_cam(side, camera) + if base_T_cam is not None: + out[side] = base_T_cam + return out + + def to_jsonable(self) -> dict[str, Any]: + """Serialize this calibration to plain values accepted by Zarr attributes.""" + out: dict[str, Any] = { + "reference_frame": self.reference_frame, + "cameras": { + name: cam.to_jsonable() for name, cam in self.cameras.items() + }, + } + if self.arm_bases: + out["arm_bases"] = { + side: np.asarray(T, dtype=np.float64).tolist() + for side, T in self.arm_bases.items() + } + return out + + +def camera_name(image_key: str) -> str | None: + """Extract the camera name after the final ``images.`` in an array key. + + Args: + image_key: An array key such as ``"images.front_1"``. + + Returns: + The camera name, or ``None`` if the key names no image stream. + """ + marker = IMAGE_KEY_PREFIX + index = image_key.rfind(marker) + if index == -1: + return None + return image_key[index + len(marker) :] or None + + +def uncalibrated_cameras(image_keys, calibration: Calibration | None) -> list[str]: + """Return unique image-stream camera names without a declared ``K``. + + Args: + image_keys: The episode's image array keys. + calibration: The episode calibration, or ``None``. + + Returns: + The camera names, in the order the keys give them and without + duplicates, that carry no ``K``. + """ + missing: list[str] = [] + seen: set[str] = set() + for key in image_keys: + name = camera_name(key) + if name is None or name in seen: + continue + seen.add(name) + if calibration is None or calibration.K(name) is None: + missing.append(name) + return missing + + +def _parse_reference_frame(value: Any, cameras, where: str) -> str: + if not isinstance(value, str) or not value: + raise CalibrationError(f"{where}: `reference_frame` must be a non-empty string") + if value in STATIC_REFERENCE_FRAMES: + return value + if value.startswith(CAMERA_FRAME_PREFIX): + camera = value[len(CAMERA_FRAME_PREFIX) :] + if camera and camera in cameras: + return value + raise CalibrationError( + f"{where}: reference_frame {value!r} names a camera that this episode " + f"does not declare; declared cameras are {sorted(cameras)}" + ) + raise CalibrationError( + f"{where}: unknown reference_frame {value!r}; expected one of " + f"{sorted(STATIC_REFERENCE_FRAMES)} or 'camera:'" + ) + + +def _parse_distortion(raw: Any, model: str, where: str) -> tuple[float, ...]: + """Return a finite distortion vector whose length is valid for ``model``.""" + if raw is None: + raw = () + if isinstance(raw, (str, bytes)) or not isinstance(raw, (list, tuple, np.ndarray)): + raise CalibrationError( + f"{where}.distortion: expected a list of coefficients, got {raw!r}" + ) + try: + coefficients = tuple(float(c) for c in raw) + except (TypeError, ValueError) as exc: + raise CalibrationError( + f"{where}.distortion: contains a non-numeric coefficient ({exc})" + ) from exc + if not all(np.isfinite(coefficients)): + raise CalibrationError(f"{where}.distortion: contains a non-finite value") + allowed = CAMERA_MODELS[model] + if len(coefficients) not in allowed: + raise CalibrationError( + f"{where}.distortion: model {model} takes " + f"{sorted(allowed)} coefficients, got {len(coefficients)}" + ) + return coefficients + + +def _parse_camera(name: str, block: Any, where: str) -> CameraCalibration: + if not isinstance(block, Mapping): + raise CalibrationError(f"{where}: must be a mapping, got {block!r}") + unknown = sorted(set(block) - _CAMERA_FIELDS) + if unknown: + raise CalibrationError( + f"{where}: unknown field(s) {unknown}; expected some of " + f"{sorted(_CAMERA_FIELDS)}" + ) + + K = block.get("K") + + model = block.get("model", DEFAULT_CAMERA_MODEL) + if model not in CAMERA_MODELS: + raise CalibrationError( + f"{where}.model: unknown camera model {model!r}; expected one of " + f"{sorted(CAMERA_MODELS)}" + ) + distortion = _parse_distortion(block.get("distortion"), model, where) + + resolution = block.get("resolution") + if resolution is not None: + if ( + not isinstance(resolution, (list, tuple)) + or len(resolution) != 2 + or not all(isinstance(v, int) and v > 0 for v in resolution) + ): + raise CalibrationError( + f"{where}.resolution: expected [width, height] positive ints, " + f"got {resolution!r}" + ) + resolution = (int(resolution[0]), int(resolution[1])) + + rectified = block.get("rectified", True) + if not isinstance(rectified, bool): + raise CalibrationError( + f"{where}.rectified: expected a boolean, got {rectified!r}" + ) + + ref_T_cam = block.get("ref_T_cam") + return CameraCalibration( + name=name, + K=None if K is None else _camera_matrix(K, f"{where}.K"), + model=model, + distortion=distortion, + resolution=resolution, + rectified=rectified, + ref_T_cam=( + None + if ref_T_cam is None + else _matrix(ref_T_cam, (4, 4), f"{where}.ref_T_cam") + ), + ) + + +def parse_calibration(block: Any, where: str = "calibration") -> Calibration: + """Validate and normalize one ``calibration`` attribute block. + + Args: + block: The mapping stored under ``zarr.attrs["calibration"]``. + where: A label used in error messages. + + Returns: + The parsed calibration. + + Raises: + CalibrationError: If a field is unknown, missing, or malformed. + """ + if isinstance(block, Calibration): + return block + if not isinstance(block, Mapping): + raise CalibrationError(f"{where}: must be a mapping, got {block!r}") + unknown = sorted(set(block) - _CALIBRATION_FIELDS) + if unknown: + raise CalibrationError( + f"{where}: unknown field(s) {unknown}; expected some of " + f"{sorted(_CALIBRATION_FIELDS)}" + ) + + raw_cameras = block.get("cameras") + if not isinstance(raw_cameras, Mapping) or not raw_cameras: + raise CalibrationError( + f"{where}.cameras: expected a non-empty mapping from camera name to " + f"its calibration, got {raw_cameras!r}" + ) + cameras = { + name: _parse_camera(name, cam, f"{where}.cameras[{name!r}]") + for name, cam in raw_cameras.items() + } + + if "reference_frame" not in block: + raise CalibrationError(f"{where}: missing required field 'reference_frame'") + reference_frame = _parse_reference_frame( + block["reference_frame"], cameras, where + ) + + raw_arm_bases = block.get("arm_bases") or {} + if not isinstance(raw_arm_bases, Mapping): + raise CalibrationError( + f"{where}.arm_bases: expected a mapping from side to a 4x4 " + f"ref_T_armbase matrix, got {raw_arm_bases!r}" + ) + arm_bases = { + side: _matrix(T, (4, 4), f"{where}.arm_bases[{side!r}]") + for side, T in raw_arm_bases.items() + } + + return Calibration( + reference_frame=reference_frame, cameras=cameras, arm_bases=arm_bases + ) + + +def lift_legacy_calibration( + intrinsics: Any = None, extrinsics: Any = None +) -> Calibration | None: + """Build a :class:`Calibration` from the legacy attribute pair. + + ``intrinsics`` maps camera names to their ``K`` matrices. A bare matrix is + assigned to ``front_1``. ``extrinsics`` maps each arm side to + ``base_T_cam`` for ``front_1``; lifting therefore uses + ``camera:front_1`` as the reference frame and stores the inverse matrix as + ``arm_bases[side]``. + + Args: + intrinsics: The ``intrinsics`` attribute, or ``None``. + extrinsics: The ``extrinsics`` attribute, or ``None``. + + Returns: + The lifted calibration, or ``None`` if both attributes are absent. + + Raises: + CalibrationError: If either attribute is malformed. + """ + if not intrinsics and not extrinsics: + return None + + if isinstance(intrinsics, Mapping): + raw_cameras = list(intrinsics.items()) + elif intrinsics is not None: + raw_cameras = [(LEGACY_REFERENCE_CAMERA, intrinsics)] + else: + raw_cameras = [] + + cameras = { + str(name): CameraCalibration( + name=str(name), K=_camera_matrix(K, f"intrinsics[{name!r}]") + ) + for name, K in raw_cameras + } + + arm_bases: dict[str, np.ndarray] = {} + if extrinsics is not None: + if not isinstance(extrinsics, Mapping): + raise CalibrationError( + f"extrinsics: expected a mapping from arm side to a 4x4 " + f"base_T_cam matrix, got {extrinsics!r}" + ) + for side, base_T_cam in extrinsics.items(): + matrix = _matrix(base_T_cam, (4, 4), f"extrinsics[{side!r}]") + arm_bases[str(side)] = np.linalg.inv(matrix) + + # Legacy extrinsics are poses of ``front_1`` in each arm base. Declare that + # camera even without intrinsics so ``base_T_cam`` remains composable. + if arm_bases and LEGACY_REFERENCE_CAMERA not in cameras: + cameras[LEGACY_REFERENCE_CAMERA] = CameraCalibration( + name=LEGACY_REFERENCE_CAMERA + ) + + reference_frame = ( + f"{CAMERA_FRAME_PREFIX}{LEGACY_REFERENCE_CAMERA}" + if LEGACY_REFERENCE_CAMERA in cameras + else f"{CAMERA_FRAME_PREFIX}{next(iter(cameras))}" + ) + return Calibration( + reference_frame=reference_frame, + cameras=cameras, + arm_bases=arm_bases, + legacy=True, + ) + + +def read_calibration(attrs: Mapping[str, Any]) -> Calibration | None: + """Parse current or legacy calibration from one episode's attributes. + + A non-empty ``calibration`` block takes precedence. Otherwise the legacy + ``intrinsics`` and ``extrinsics`` attributes are lifted into the same + representation. + + Args: + attrs: The episode's Zarr attributes. + + Returns: + The episode calibration, or ``None`` if the episode states none. + + Raises: + CalibrationError: If the stored calibration is malformed. + """ + block = attrs.get("calibration") + if block: + return parse_calibration(block) + return lift_legacy_calibration( + attrs.get("intrinsics"), attrs.get("extrinsics") + ) + + +__all__ = [ + "CAMERA_FRAME_PREFIX", + "CAMERA_MODELS", + "DEFAULT_CAMERA_MODEL", + "IMAGE_KEY_PREFIX", + "LEGACY_REFERENCE_CAMERA", + "STATIC_REFERENCE_FRAMES", + "Calibration", + "CalibrationError", + "CameraCalibration", + "camera_name", + "lift_legacy_calibration", + "parse_calibration", + "read_calibration", + "uncalibrated_cameras", +] diff --git a/egomimic/rldb/zarr/episode_attrs.py b/egomimic/rldb/zarr/episode_attrs.py new file mode 100644 index 000000000..5bcbd916f --- /dev/null +++ b/egomimic/rldb/zarr/episode_attrs.py @@ -0,0 +1,44 @@ +"""Shared names and normalization helpers for episode attributes. + +The writer and readers depend on this module without depending on each other. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +#: A finished recording eligible for staging and resolver-based loading. +DATA_STATUS_COMPLETE = "complete" + +#: A delivery-format example that staging and dataset resolvers must exclude. +DATA_STATUS_STRUCTURAL_SAMPLE = "structural_sample" + +DATA_STATUS_VALUES = (DATA_STATUS_COMPLETE, DATA_STATUS_STRUCTURAL_SAMPLE) + + +def data_status(attrs: Mapping) -> str: + """Return the stored status, defaulting a missing or falsey value to complete. + + The default preserves the behavior of episodes written before + ``data_status`` was introduced. This helper does not validate a non-empty + stored value; the writer and schema validator do that separately. + + Args: + attrs: The episode's Zarr attributes. + """ + value = attrs.get("data_status") + return DATA_STATUS_COMPLETE if not value else str(value) + + +def is_complete(attrs: Mapping) -> bool: + """Return whether the normalized status is exactly ``complete``.""" + return data_status(attrs) == DATA_STATUS_COMPLETE + + +__all__ = [ + "DATA_STATUS_COMPLETE", + "DATA_STATUS_STRUCTURAL_SAMPLE", + "DATA_STATUS_VALUES", + "data_status", + "is_complete", +] diff --git a/egomimic/rldb/zarr/schema/episode_v3.yaml b/egomimic/rldb/zarr/schema/episode_v3.yaml new file mode 100644 index 000000000..786a57e6d --- /dev/null +++ b/egomimic/rldb/zarr/schema/episode_v3.yaml @@ -0,0 +1,278 @@ +# The declarative portion of the episode contract. +# +# This file selects fields and named checks, supplies their parameters, and +# assigns each rule a severity class. `validate.py` implements the schema +# interpreter and the named check predicates referenced here. +# +# Run it with: +# python -m egomimic.rldb.zarr.validate [RULE FLAGS] +# +# --------------------------------------------------------------------------- +# attributes +# required true | false. A missing required attribute or array emits a +# finding. An optional field emits no finding when absent. +# severity integrity | adoption | coverage +# Integrity findings are always errors. Adoption and coverage +# findings are errors unless their named rule is waived for this +# invocation. Present fields must meet their declared type and +# shape even when their presence requirement is waived. +# name Stable rule name when the attribute or key is not the desired +# CLI flag stem. Dashes replace underscores in CLI flags. +# why User-facing reason for a waivable rule. This is required because +# the schema generates CLI help and must explain what accepting a +# gap gives up without relying on a separate rollout document. +# type str | int | bool | mapping +# min smallest accepted value for an int +# choices the accepted values +# check an attribute predicate implemented in `validate.py` +# +# checks +# name a predicate implemented in `validate.py` +# required whether a predicate failure emits a finding +# severity controls whether that finding can be waived, as described above +# predicate-specific parameters +# +# arrays +# key an array key. `{side}` expands over arity-derived candidates, +# filtered by resolved end-effectors; `*` selects stored matches. +# required true | false, as above. A wildcard key is never required. +# severity controls missing-key severity. Shape and dtype failures are +# integrity errors for every array, including a waived one. +# when conditions on the resolved embodiment; skip the rule when any +# condition does not hold +# shape expected dimensions, see below +# dtype int | float | object | any +# +# shape dimensions +# T `total_frames`. Axis 0 may be longer, because a stored array may +# carry a padded tail; every other axis is exact. +# arm_dof the platform's joints per arm +# aux_dof the platform's auxiliary chain +# ee_dof the side's end-effector joints +# kp3 three coordinates per slot in the side's keypoint topology +# exactly that many +# "*" any +# +# when conditions +# platform_kind robot | human +# has_arm_chain the platform declares `arm_dof` +# has_aux_chain the platform declares `aux` +# end_effector_class parallel_jaw | human_hand | dexterous_hand +# --------------------------------------------------------------------------- + +version: v3.1 +known_versions: [v3.0, v3.1] + +attributes: + embodiment: + required: true + severity: integrity + type: str + check: embodiment_name + total_frames: + required: true + severity: integrity + type: int + min: 1 + fps: + required: true + severity: integrity + type: int + min: 1 + data_status: + required: true + severity: adoption + why: >- + Require an explicit data status so an incomplete structural sample + cannot inherit the legacy "complete" default and enter training. + type: str + choices: [complete, structural_sample] + task_name: + required: true + severity: integrity + type: str + task_description: + required: false + severity: integrity + type: str + features: + required: true + severity: integrity + type: mapping + schema_version: + required: true + severity: adoption + why: >- + Require a schema version so readers can select the episode contract + explicitly instead of inferring it from whichever fields are present. + type: str + check: schema_version + morphology: + required: false + severity: integrity + type: mapping + check: morphology + calibration: + name: calibration_block + required: true + severity: adoption + why: >- + Require the calibration block so camera models, distortion, resolution, + and frame relationships are explicit; waive only when legacy calibration + attributes contain enough information for the compatibility reader. + type: mapping + check: calibration + intrinsics: + required: false + severity: integrity + type: mapping + extrinsics: + required: false + severity: integrity + type: mapping + +checks: + # Require at least one camera matrix, read from the current `calibration` + # block or lifted from the legacy attributes. + - name: calibration_present + required: true + severity: integrity + + # Match each stored `images.` stream to a camera entry carrying K. + - name: camera_coverage + required: true + severity: coverage + why: >- + Require a camera matrix for every image stream because projection and + reprojection checks cannot interpret an uncalibrated view; waive only + after explicitly accepting that physical calibration gap. + + # For a listed 7D pose track with at least two retained rows, reject identical + # rows or more than 1% exact +identity/-identity quaternions. + - name: pose_degeneracy + required: true + severity: integrity + suffixes: [obs_ee_pose, cmd_ee_pose, obs_head_pose] + identity_rotation_fraction: 0.01 + + # Reject an exact identity pose stored for a non-reference camera, or a + # numerically identity `base_T_cam` composed for an arm and the default camera. + - name: calibration_degeneracy + required: true + severity: integrity + + # Reject the exact synthetic-intrinsics signature `fx == fy == width`, + # `cx == width / 2`, and `cy == height / 2`. + - name: intrinsics_signature + required: true + severity: integrity + + # Reject non-increasing retained timestamps, stored keys containing + # `relative_timestamp`, and a power-of-two step GCD of at least 64 ns (the + # heuristic signature of float seconds converted to integer nanoseconds). + - name: timestamps + required: true + severity: integrity + key: obs_rgb_timestamps_ns + banned_key_substrings: [relative_timestamp] + + # Valid annotation intervals must jointly cover at least 90% of + # `[0, total_frames)`. + - name: annotation_coverage + required: true + severity: coverage + why: >- + Require annotations to cover at least 90% of retained frames because an + unlabelled tail gives those frames no reliable language target; waive + only after accepting the gap instead of annotating or trimming the tail. + key: annotations + minimum: 0.9 + + # Forbid the exact substring ` | ` in annotation text and task_description; + # structured metadata needs a separate schema field. + - name: annotation_text + required: true + severity: integrity + key: annotations + banned_delimiters: [" | "] + +arrays: + - key: "{side}.obs_ee_pose" + required: true + severity: integrity + shape: [T, 7] + dtype: float + # Observed end-effector-root pose. Require one for every resolved side, + # independent of the end-effector class. + + - key: "{side}.cmd_ee_pose" + required: false + severity: integrity + shape: [T, 7] + dtype: float + + - key: "{side}.obs_joints" + when: {has_arm_chain: true} + required: false + severity: integrity + shape: [T, arm_dof] + dtype: float + + - key: "{side}.cmd_joints" + when: {has_arm_chain: true} + required: false + severity: integrity + shape: [T, arm_dof] + dtype: float + + - key: "{side}.obs_gripper" + when: {end_effector_class: parallel_jaw} + required: true + severity: integrity + shape: [T, 1] + dtype: float + + - key: "{side}.cmd_gripper" + when: {end_effector_class: parallel_jaw} + required: true + severity: integrity + shape: [T, 1] + dtype: float + + - key: "{side}.obs_keypoints" + when: {end_effector_class: human_hand} + required: true + severity: integrity + shape: [T, kp3] + dtype: float + + - key: obs_head_pose + when: {platform_kind: human} + required: true + severity: integrity + shape: [T, 7] + dtype: float + + # Canonical RGB timestamp array. When present it must be one-dimensional, + # integer-valued, and at least `total_frames` long. + - key: obs_rgb_timestamps_ns + name: rgb_timestamps + required: true + severity: coverage + why: >- + Require integer UTC-nanosecond RGB timestamps so streams can be aligned + without reconstructing the capture clock after the recording is over. + shape: [T] + dtype: int + + - key: "images.*" + required: false + severity: integrity + shape: [T] + dtype: object + + - key: annotations + required: false + severity: integrity + shape: ["*"] + dtype: object diff --git a/egomimic/rldb/zarr/test_calibration.py b/egomimic/rldb/zarr/test_calibration.py new file mode 100644 index 000000000..f16304c21 --- /dev/null +++ b/egomimic/rldb/zarr/test_calibration.py @@ -0,0 +1,442 @@ +"""Test the ``calibration`` attribute block and the legacy read shim.""" + +import numpy as np +import pytest + +from egomimic.rldb.embodiment.eva import Eva, _build_eva_bimanual_transform_list +from egomimic.rldb.zarr.action_chunk_transforms import ( + ActionChunkCoordinateFrameTransform, +) +from egomimic.rldb.zarr.calibration import ( + Calibration, + CalibrationError, + CameraCalibration, + camera_name, + lift_legacy_calibration, + parse_calibration, + read_calibration, + uncalibrated_cameras, +) +from egomimic.rldb.zarr.zarr_dataset_multi import ZarrDataset, ZarrEpisode +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +K_FRONT = np.array( + [ + [200.0, 0.0, 160.0, 0.0], + [0.0, 200.0, 120.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + ] +) + + +def _calibration_block() -> dict: + return { + "reference_frame": "robot_base", + "cameras": { + "front_1": { + "K": K_FRONT.tolist(), + "resolution": [320, 240], + "rectified": True, + "ref_T_cam": np.eye(4).tolist(), + }, + "left_wrist": {"K": K_FRONT.tolist(), "rectified": False}, + }, + "arm_bases": {"left": np.eye(4).tolist()}, + } + + +def test_parse_round_trips_through_jsonable() -> None: + calibration = parse_calibration(_calibration_block()) + assert calibration.reference_frame == "robot_base" + assert calibration.reference_camera is None + assert calibration.cameras["front_1"].resolution == (320, 240) + assert calibration.cameras["left_wrist"].rectified is False + reparsed = parse_calibration(calibration.to_jsonable()) + assert reparsed.reference_frame == calibration.reference_frame + assert set(reparsed.cameras) == set(calibration.cameras) + for name, camera in calibration.cameras.items(): + other = reparsed.cameras[name] + assert (other.resolution, other.rectified) == ( + camera.resolution, + camera.rectified, + ) + np.testing.assert_array_equal(other.K, camera.K) + np.testing.assert_array_equal( + reparsed.arm_bases["left"], calibration.arm_bases["left"] + ) + + +def test_default_camera_prefers_the_front_camera() -> None: + calibration = parse_calibration( + { + "reference_frame": "robot_base", + "cameras": { + "left_wrist": {"K": K_FRONT.tolist()}, + "front_1": {"K": (2 * K_FRONT).tolist()}, + }, + } + ) + assert calibration.default_camera() == "front_1" + np.testing.assert_allclose(calibration.K(), 2 * K_FRONT) + + +def test_reference_camera_pose_is_the_identity() -> None: + calibration = parse_calibration( + { + "reference_frame": "camera:front_1", + "cameras": {"front_1": {"K": K_FRONT.tolist()}}, + } + ) + np.testing.assert_allclose(calibration.ref_T_cam("front_1"), np.eye(4)) + assert calibration.ref_T_cam("left_wrist") is None + + +@pytest.mark.parametrize( + "block, message", + [ + ({"cameras": {"front_1": {}}}, "reference_frame"), + ({"reference_frame": "robot_base", "cameras": {}}, "non-empty mapping"), + ( + {"reference_frame": "camera:missing", "cameras": {"front_1": {}}}, + "does not declare", + ), + ( + {"reference_frame": "elbow", "cameras": {"front_1": {}}}, + "unknown reference_frame", + ), + ( + {"reference_frame": "robot_base", "cameras": {"front_1": {"fx": 1}}}, + "unknown field", + ), + ( + { + "reference_frame": "robot_base", + "cameras": {"front_1": {"K": [[1.0, 2.0], [3.0, 4.0]]}}, + }, + "3x4 camera matrix", + ), + ( + { + "reference_frame": "robot_base", + "cameras": {"front_1": {}}, + "arm_bases": {"left": np.eye(3).tolist()}, + }, + r"expected shape \(4, 4\)", + ), + ], +) +def test_parse_rejects_malformed_blocks(block, message) -> None: + with pytest.raises(CalibrationError, match=message): + parse_calibration(block) + + +def test_a_camera_declares_pinhole_and_no_distortion_by_default() -> None: + calibration = parse_calibration( + { + "reference_frame": "robot_base", + "cameras": {"front_1": {"K": K_FRONT.tolist()}}, + } + ) + camera = calibration.cameras["front_1"] + assert (camera.model, camera.distortion) == ("PINHOLE", ()) + assert calibration.to_jsonable()["cameras"]["front_1"]["model"] == "PINHOLE" + + +def test_a_declared_lens_model_round_trips() -> None: + block = { + "reference_frame": "robot_base", + "cameras": { + "front_1": { + "K": K_FRONT.tolist(), + "model": "KANNALA_BRANDT", + "distortion": [0.1, -0.02, 0.003, -0.0004], + "rectified": False, + } + }, + } + camera = parse_calibration(block).cameras["front_1"] + assert camera.model == "KANNALA_BRANDT" + assert camera.distortion == (0.1, -0.02, 0.003, -0.0004) + reparsed = parse_calibration(parse_calibration(block).to_jsonable()) + assert reparsed.cameras["front_1"].distortion == camera.distortion + + +@pytest.mark.parametrize( + "camera, message", + [ + ({"model": "FISHEYE"}, "unknown camera model"), + ({"distortion": [0.1, 0.2]}, "model PINHOLE takes"), + ({"model": "KANNALA_BRANDT", "distortion": [0.1]}, "takes"), + ({"distortion": "0.1"}, "expected a list"), + ({"model": "OPENCV", "distortion": [0.1, 0.2, 0.3, "x"]}, "non-numeric"), + ], +) +def test_a_camera_model_and_its_coefficients_must_agree(camera, message) -> None: + with pytest.raises(CalibrationError, match=message): + parse_calibration( + {"reference_frame": "robot_base", "cameras": {"front_1": camera}} + ) + + +def test_a_bare_3x3_camera_matrix_is_padded() -> None: + calibration = parse_calibration( + { + "reference_frame": "robot_base", + "cameras": {"front_1": {"K": K_FRONT[:, :3].tolist()}}, + } + ) + np.testing.assert_allclose(calibration.K("front_1"), K_FRONT) + + +def test_legacy_attributes_lift_without_changing_base_T_cam() -> None: + calibration = lift_legacy_calibration( + intrinsics={"front_1": K_FRONT}, extrinsics=Eva.EXTRINSICS + ) + assert calibration.legacy + assert calibration.reference_frame == "camera:front_1" + np.testing.assert_allclose(calibration.K(), K_FRONT) + for side, base_T_cam in Eva.EXTRINSICS.items(): + # `arm_bases` holds ref_T_armbase, the inverse of the stored matrix. + np.testing.assert_allclose( + calibration.arm_bases[side], np.linalg.inv(base_T_cam), atol=1e-12 + ) + np.testing.assert_allclose( + calibration.base_T_cam(side), base_T_cam, atol=1e-12 + ) + + +def test_legacy_extrinsics_without_intrinsics_still_compose() -> None: + calibration = lift_legacy_calibration(extrinsics=Eva.EXTRINSICS) + assert calibration.K() is None + np.testing.assert_allclose( + calibration.base_T_cam("left"), Eva.EXTRINSICS["left"], atol=1e-12 + ) + + +def test_read_calibration_prefers_the_block_over_the_legacy_pair() -> None: + attrs = { + "calibration": _calibration_block(), + "intrinsics": {"front_1": (7 * K_FRONT).tolist()}, + } + calibration = read_calibration(attrs) + assert calibration is not None and not calibration.legacy + np.testing.assert_allclose(calibration.K(), K_FRONT) + assert read_calibration({}) is None + + +def test_camera_name_reads_the_stream_out_of_an_array_key() -> None: + assert camera_name("images.front_1") == "front_1" + assert camera_name("observations.images.left_wrist") == "left_wrist" + assert camera_name("left.obs_ee_pose") is None + assert camera_name("images.") is None + + +def test_uncalibrated_cameras_lists_streams_without_a_matrix() -> None: + calibration = Calibration( + reference_frame="robot_base", + cameras={ + "front_1": CameraCalibration(name="front_1", K=K_FRONT), + "left_wrist": CameraCalibration(name="left_wrist"), + }, + ) + keys = ["images.front_1", "images.left_wrist", "images.right_wrist"] + assert uncalibrated_cameras(keys, calibration) == ["left_wrist", "right_wrist"] + assert uncalibrated_cameras(keys, None) == [ + "front_1", + "left_wrist", + "right_wrist", + ] + assert uncalibrated_cameras(["left.obs_gripper"], None) == [] + + +def test_writer_camera_coverage_can_be_required_explicitly(tmp_path, caplog) -> None: + images = { + "images.front_1": np.zeros((4, 8, 8, 3), dtype=np.uint8), + "images.left_wrist": np.zeros((4, 8, 8, 3), dtype=np.uint8), + } + with caplog.at_level("WARNING"): + _write_episode( + tmp_path / "partial.zarr", + image_data=images, + intrinsics={"front_1": K_FRONT}, + ) + assert "left_wrist" in caplog.text + + with pytest.raises(ValueError, match=r"no camera matrix.*left_wrist"): + _write_episode( + tmp_path / "required.zarr", + image_data=images, + intrinsics={"front_1": K_FRONT}, + require_camera_coverage=True, + ) + + +def test_full_coverage_passes_when_camera_coverage_is_required( + tmp_path, caplog +) -> None: + with caplog.at_level("WARNING"): + _write_episode( + tmp_path / "covered.zarr", + image_data={"images.front_1": np.zeros((4, 8, 8, 3), dtype=np.uint8)}, + calibration={ + "reference_frame": "camera:front_1", + "cameras": {"front_1": {"K": K_FRONT.tolist()}}, + }, + require_camera_coverage=True, + ) + assert "no camera matrix" not in caplog.text + + +def _write_episode(episode_path, **kwargs) -> None: + ZarrWriter.create_and_write( + episode_path=episode_path, + numeric_data={"left.obs_gripper": np.zeros((4, 1))}, + embodiment="eva_bimanual", + chunk_timesteps=4, + **kwargs, + ) + + +def test_the_writer_stores_the_block_and_the_legacy_pair(tmp_path) -> None: + episode_path = tmp_path / "calibrated.zarr" + _write_episode(episode_path, calibration=_calibration_block()) + + episode = ZarrEpisode(episode_path) + assert episode.metadata["calibration"]["reference_frame"] == "robot_base" + # Readers that predate the block still find the pair they expect. + np.testing.assert_allclose( + np.asarray(episode.metadata["intrinsics"]["front_1"]), K_FRONT + ) + np.testing.assert_allclose( + np.asarray(episode.metadata["extrinsics"]["left"]), np.eye(4) + ) + np.testing.assert_allclose(episode.calibration.K(), K_FRONT) + + +def test_a_legacy_episode_reads_through_the_shim(tmp_path) -> None: + episode_path = tmp_path / "legacy.zarr" + _write_episode( + episode_path, intrinsics={"front_1": K_FRONT}, extrinsics=Eva.EXTRINSICS + ) + + episode = ZarrEpisode(episode_path) + assert "calibration" not in episode.metadata + calibration = episode.calibration + assert calibration is not None and calibration.legacy + np.testing.assert_allclose(calibration.K(), K_FRONT) + np.testing.assert_allclose( + calibration.base_T_cam("right"), Eva.EXTRINSICS["right"], atol=1e-12 + ) + + +def test_the_writer_rejects_an_episode_with_no_camera_matrix(tmp_path) -> None: + with pytest.raises(ValueError, match="Camera intrinsics"): + _write_episode(tmp_path / "blank.zarr") + + +# -------------------------------------------------------------------------- +# The training path reads the rig the episode declares +# -------------------------------------------------------------------------- + +_EVA_KEY_MAP = { + "left.obs_ee_pose": {"zarr_key": "left.obs_ee_pose"}, + "right.obs_ee_pose": {"zarr_key": "right.obs_ee_pose"}, + "left.obs_gripper": {"zarr_key": "left.obs_gripper"}, + "right.obs_gripper": {"zarr_key": "right.obs_gripper"}, + "left.cmd_ee_pose": {"zarr_key": "left.cmd_ee_pose", "horizon": 4}, + "right.cmd_ee_pose": {"zarr_key": "right.cmd_ee_pose", "horizon": 4}, + "left.cmd_gripper": {"zarr_key": "left.cmd_gripper", "horizon": 4}, + "right.cmd_gripper": {"zarr_key": "right.cmd_gripper", "horizon": 4}, +} + + +def _pose_sequence(length: int, *, x_offset: float = 0.0) -> np.ndarray: + poses = np.zeros((length, 7)) + poses[:, 0] = x_offset + np.arange(length) * 0.01 + poses[:, 3] = 1.0 + return poses + + +def _eva_numeric_data(length: int) -> dict: + return { + "left.obs_ee_pose": _pose_sequence(length), + "right.obs_ee_pose": _pose_sequence(length, x_offset=0.2), + "left.obs_gripper": np.full((length, 1), 0.25), + "right.obs_gripper": np.full((length, 1), 0.75), + "left.cmd_ee_pose": _pose_sequence(length, x_offset=0.1), + "right.cmd_ee_pose": _pose_sequence(length, x_offset=0.3), + "left.cmd_gripper": np.linspace(0.0, 1.0, length)[:, None], + "right.cmd_gripper": np.linspace(1.0, 0.0, length)[:, None], + } + + +def _other_rig() -> dict: + """Translate each fallback camera pose by ``[0.05, -0.11, 0.07]`` metres.""" + rig = {} + for side, base_T_cam in Eva.EXTRINSICS.items(): + moved = base_T_cam.copy() + moved[:3, 3] += np.array([0.05, -0.11, 0.07]) + rig[side] = moved + return rig + + +def _eva_actions(episode_path, extrinsics) -> np.ndarray: + """Write one EVA episode and return its transformed action chunk.""" + ZarrWriter.create_and_write( + episode_path=episode_path, + numeric_data=_eva_numeric_data(5), + embodiment="eva_bimanual", + intrinsics={"front_1": K_FRONT}, + extrinsics=extrinsics, + chunk_timesteps=4, + ) + dataset = ZarrDataset( + Episode_path=episode_path, + key_map=dict(_EVA_KEY_MAP), + transform_list=_build_eva_bimanual_transform_list( + chunk_length=6, stride=1, is_quat=True + ), + ) + return dataset[1]["actions_cartesian"].numpy() + + +def test_reading_the_episode_is_a_no_op_on_the_current_corpus(tmp_path) -> None: + stored = _eva_actions(tmp_path / "stored.zarr", Eva.EXTRINSICS) + absent = _eva_actions(tmp_path / "absent.zarr", None) + # Storing the class-default rig and omitting it must exercise the same + # transform because omission selects that fallback. + np.testing.assert_allclose(stored, absent, atol=1e-9) + + +def test_the_pipeline_follows_the_rig_the_episode_declares(tmp_path) -> None: + on_eva_rig = _eva_actions(tmp_path / "eva_rig.zarr", Eva.EXTRINSICS) + on_other_rig = _eva_actions(tmp_path / "other_rig.zarr", _other_rig()) + assert not np.allclose(on_eva_rig, on_other_rig, atol=1e-6) + + +def test_a_second_rig_matches_what_that_rig_predicts(tmp_path, monkeypatch) -> None: + other = _other_rig() + declared = _eva_actions(tmp_path / "declared.zarr", other) + + # An episode without extrinsics falls back to the class constant, so + # pointing the constant at the second rig must reproduce the same numbers. + monkeypatch.setattr(Eva, "EXTRINSICS", other) + fallback = _eva_actions(tmp_path / "fallback.zarr", None) + np.testing.assert_allclose(declared, fallback, atol=1e-9) + + +def test_a_transform_fallback_never_overrides_a_batch_value() -> None: + identity_pose = np.array([1.0, 2.0, 3.0, 1.0, 0.0, 0.0, 0.0]) + transform = ActionChunkCoordinateFrameTransform( + target_world="target", + chunk_world="chunk", + transformed_key_name="out", + extra_batch_key={"target": np.array([9.0, 9.0, 9.0, 1.0, 0.0, 0.0, 0.0])}, + mode="xyzwxyz", + ) + batch = {"target": identity_pose, "chunk": np.tile(identity_pose, (2, 1))} + + out = transform.transform(batch) + + np.testing.assert_array_equal(out["target"], identity_pose) + np.testing.assert_allclose(out["out"][:, :3], 0.0, atol=1e-12) diff --git a/egomimic/rldb/zarr/test_data_status.py b/egomimic/rldb/zarr/test_data_status.py new file mode 100644 index 000000000..285eb3426 --- /dev/null +++ b/egomimic/rldb/zarr/test_data_status.py @@ -0,0 +1,116 @@ +"""Test status writing, legacy defaults, validation, and resolver filtering.""" + +import numpy as np +import pytest +import zarr + +from egomimic.rldb.zarr.episode_attrs import ( + DATA_STATUS_COMPLETE, + DATA_STATUS_STRUCTURAL_SAMPLE, + data_status, + is_complete, +) +from egomimic.rldb.zarr.validate import ERROR, WARNING, validate_episode +from egomimic.rldb.zarr.zarr_dataset_multi import ( + EpisodeResolver, + ZarrDataset, + ZarrEpisode, +) +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +K = np.array([[200.0, 0.0, 160.0, 0.0], [0.0, 200.0, 120.0, 0.0], [0.0, 0.0, 1.0, 0.0]]) +LENGTH = 4 + + +def _write(path, **kwargs) -> None: + ZarrWriter.create_and_write( + episode_path=path, + numeric_data={"left.obs_gripper": np.zeros((LENGTH, 1))}, + embodiment="eva_bimanual", + chunk_timesteps=LENGTH, + intrinsics={"front_1": K}, + **kwargs, + ) + + +def test_an_episode_without_the_attribute_reads_as_complete() -> None: + assert data_status({}) == DATA_STATUS_COMPLETE + assert is_complete({}) + assert not is_complete({"data_status": DATA_STATUS_STRUCTURAL_SAMPLE}) + + +def test_the_writer_records_the_status(tmp_path) -> None: + _write(tmp_path / "complete.zarr") + _write(tmp_path / "sample.zarr", data_status=DATA_STATUS_STRUCTURAL_SAMPLE) + + assert ZarrEpisode(tmp_path / "complete.zarr").data_status == DATA_STATUS_COMPLETE + assert ( + ZarrEpisode(tmp_path / "sample.zarr").data_status + == DATA_STATUS_STRUCTURAL_SAMPLE + ) + + +def test_the_writer_rejects_an_unknown_status(tmp_path) -> None: + with pytest.raises(ValueError, match="data_status must be one of"): + _write(tmp_path / "bad.zarr", data_status="mostly_done") + + +def test_the_resolver_refuses_a_structural_sample(tmp_path, caplog) -> None: + _write(tmp_path / "complete.zarr") + _write(tmp_path / "sample.zarr", data_status=DATA_STATUS_STRUCTURAL_SAMPLE) + resolver = EpisodeResolver(tmp_path, key_map={}) + + with caplog.at_level("WARNING"): + datasets = resolver._load_zarr_datasets( + search_path=tmp_path, valid_folder_names={"complete", "sample"} + ) + + assert set(datasets) == {"complete"} + assert "not 'complete'" in caplog.text + + +def test_the_dataset_exposes_the_status(tmp_path) -> None: + _write(tmp_path / "sample.zarr", data_status=DATA_STATUS_STRUCTURAL_SAMPLE) + + dataset = ZarrDataset(Episode_path=tmp_path / "sample.zarr", key_map={}) + + assert dataset.data_status == DATA_STATUS_STRUCTURAL_SAMPLE + + +def test_the_validator_reads_and_checks_the_status(tmp_path) -> None: + _write(tmp_path / "sample.zarr", data_status=DATA_STATUS_STRUCTURAL_SAMPLE) + report = validate_episode(tmp_path / "sample.zarr") + assert next( + f.level for f in report.findings if f.check == "attrs.data_status" + ) == "ok" + + _write(tmp_path / "legacy.zarr") + store = zarr.open_group(str(tmp_path / "legacy.zarr"), mode="a") + del store.attrs["data_status"] + required = validate_episode(tmp_path / "legacy.zarr") + waived = validate_episode( + tmp_path / "legacy.zarr", requirements={"data_status": False} + ) + statuses = { + "required": next( + f.level for f in required.findings if f.check == "attrs.data_status" + ), + "waived": next( + f.level for f in waived.findings if f.check == "attrs.data_status" + ), + } + assert statuses == {"required": ERROR, "waived": WARNING} + + +def test_the_validator_rejects_an_unknown_status(tmp_path) -> None: + _write(tmp_path / "odd.zarr") + store = zarr.open_group(str(tmp_path / "odd.zarr"), mode="a") + store.attrs["data_status"] = "in_review" + + report = validate_episode( + tmp_path / "odd.zarr", requirements={"data_status": False} + ) + + finding = next(f for f in report.findings if f.check == "attrs.data_status") + assert finding.level == ERROR + assert "expected one of" in finding.message diff --git a/egomimic/rldb/zarr/test_validate.py b/egomimic/rldb/zarr/test_validate.py new file mode 100644 index 000000000..5f5761eb1 --- /dev/null +++ b/egomimic/rldb/zarr/test_validate.py @@ -0,0 +1,537 @@ +"""Test the schema-driven episode validator.""" + +import numpy as np +import pytest + +from egomimic.rldb.embodiment.eva import Eva +from egomimic.rldb.zarr.validate import ( + ERROR, + OK, + WARNING, + Report, + SchemaError, + _build_parser, + _level, + load_schema, + main, + validate_episode, + waivable_rules, +) +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +K = np.array([[200.0, 0.0, 160.0, 0.0], [0.0, 200.0, 120.0, 0.0], [0.0, 0.0, 1.0, 0.0]]) +LENGTH = 4 + + +def _levels(report) -> dict[str, str]: + return {f.check: f.level for f in report.findings} + + +def _poses(x_offset: float = 0.0) -> np.ndarray: + """Return a moving ``[x, y, z, qw, qx, qy, qz]`` pose track.""" + angles = np.linspace(0.1, 0.4, LENGTH) + poses = np.zeros((LENGTH, 7)) + poses[:, 0] = x_offset + np.arange(LENGTH) * 0.01 + poses[:, 3] = np.cos(angles / 2) + poses[:, 6] = np.sin(angles / 2) + return poses + + +def _eva_numeric() -> dict: + data = {} + for side, offset in (("left", 0.0), ("right", 0.2)): + data[f"{side}.obs_ee_pose"] = _poses(offset) + data[f"{side}.cmd_ee_pose"] = _poses(offset + 0.05) + data[f"{side}.obs_joints"] = np.zeros((LENGTH, 6)) + data[f"{side}.cmd_joints"] = np.zeros((LENGTH, 6)) + data[f"{side}.obs_gripper"] = np.zeros((LENGTH, 1)) + data[f"{side}.cmd_gripper"] = np.zeros((LENGTH, 1)) + data["obs_rgb_timestamps_ns"] = np.arange(LENGTH, dtype=np.int64) + return data + + +def _write_eva(path, *, numeric=None, images=None, **kwargs) -> None: + kwargs.setdefault("calibration", { + "reference_frame": "camera:front_1", + "cameras": {"front_1": {"K": K.tolist()}}, + "arm_bases": { + side: np.linalg.inv(T).tolist() for side, T in Eva.EXTRINSICS.items() + }, + }) + kwargs.setdefault("metadata_override", {"schema_version": "v3.1"}) + kwargs.setdefault("annotations", [("fold the towel", 0, LENGTH)]) + ZarrWriter.create_and_write( + episode_path=path, + numeric_data=_eva_numeric() if numeric is None else numeric, + image_data=images, + embodiment=kwargs.pop("embodiment", "eva_bimanual"), + chunk_timesteps=kwargs.pop("chunk_timesteps", LENGTH), + **kwargs, + ) + + +def test_a_complete_eva_episode_passes_the_default_contract(tmp_path) -> None: + path = tmp_path / "eva.zarr" + _write_eva(path, images={"images.front_1": np.zeros((LENGTH, 8, 8, 3), np.uint8)}) + + report = validate_episode(path) + + assert report.ok, report.text() + assert not report.warnings + assert _levels(report)["embodiment"] == OK + + +def test_a_missing_required_array_is_an_error(tmp_path) -> None: + numeric = _eva_numeric() + del numeric["right.obs_ee_pose"] + _write_eva(tmp_path / "eva.zarr", numeric=numeric) + + report = validate_episode(tmp_path / "eva.zarr") + + assert _levels(report)["right.obs_ee_pose"] == ERROR + assert not report.ok + + +def test_a_wrong_width_is_reported_with_the_dimension_that_set_it(tmp_path) -> None: + numeric = _eva_numeric() + numeric["left.obs_joints"] = np.zeros((LENGTH, 5)) + _write_eva(tmp_path / "eva.zarr", numeric=numeric) + + report = validate_episode(tmp_path / "eva.zarr") + + message = next( + f.message for f in report.findings if f.check == "left.obs_joints" + ) + assert "axis 1 is 5, expected 6 (arm_dof)" in message + + +def test_a_padded_tail_is_not_an_error(tmp_path) -> None: + path = tmp_path / "eva.zarr" + _write_eva(path, chunk_timesteps=3) # 4 frames pad out to 6 + + report = validate_episode(path) + + assert _levels(report)["left.obs_ee_pose"] == OK + assert report.ok, report.text() + + +def test_an_array_shorter_than_total_frames_is_an_error(tmp_path) -> None: + path = tmp_path / "eva.zarr" + _write_eva(path) + store = __import__("zarr").open_group(str(path), mode="a") + store.attrs["total_frames"] = LENGTH + 10 + + report = validate_episode(path) + + assert "holds 4 frames for total_frames 14" in report.text() + + +def test_each_waiver_lowers_only_its_named_rule_to_a_warning(tmp_path) -> None: + path = tmp_path / "legacy.zarr" + ZarrWriter.create_and_write( + episode_path=path, + numeric_data=_eva_numeric(), + image_data={ + "images.front_1": np.zeros((LENGTH, 8, 8, 3), np.uint8), + "images.left_wrist": np.zeros((LENGTH, 8, 8, 3), np.uint8), + }, + embodiment="eva_bimanual", + chunk_timesteps=LENGTH, + annotations=[("fold the towel", 0, LENGTH)], + intrinsics={"front_1": K}, + extrinsics=Eva.EXTRINSICS, + ) + + default = validate_episode(path) + waived = validate_episode( + path, + requirements={ + "camera_coverage": False, + "calibration_block": False, + "schema_version": False, + }, + ) + + assert not default.ok + assert _levels(default)["camera_coverage"] == ERROR + assert _levels(default)["attrs.calibration"] == ERROR + assert waived.ok + assert _levels(waived)["camera_coverage"] == WARNING + assert _levels(waived)["attrs.calibration"] == WARNING + assert "left_wrist" in waived.text() + + +def test_a_single_arm_episode_owes_only_its_own_arm(tmp_path) -> None: + numeric = { + k: v for k, v in _eva_numeric().items() if not k.startswith("right.") + } + _write_eva(tmp_path / "left.zarr", numeric=numeric, embodiment="eva_left_arm") + + report = validate_episode(tmp_path / "left.zarr") + + assert report.ok, report.text() + assert "right.obs_ee_pose" not in _levels(report) + + +def test_a_human_episode_owes_keypoints_and_a_head_pose(tmp_path) -> None: + path = tmp_path / "human.zarr" + numeric = { + "left.obs_ee_pose": _poses(), + "right.obs_ee_pose": _poses(0.2), + "left.obs_keypoints": np.zeros((LENGTH, 63)), + "obs_head_pose": _poses(), + "obs_rgb_timestamps_ns": np.arange(LENGTH, dtype=np.int64), + } + ZarrWriter.create_and_write( + episode_path=path, + numeric_data=numeric, + embodiment="human_bimanual", + chunk_timesteps=LENGTH, + annotations=[("wave", 0, LENGTH)], + intrinsics={"front_1": K}, + metadata_override={"schema_version": "v3.1"}, + ) + + report = validate_episode(path) + + levels = _levels(report) + assert levels["left.obs_keypoints"] == OK + assert levels["obs_head_pose"] == OK + # The registry declares 21 MANO slots with three coordinates per slot. + assert levels["right.obs_keypoints"] == ERROR + # The human platform has neither an arm chain nor a parallel jaw, so those + # conditional rules do not run. + assert "left.obs_gripper" not in levels + assert "left.obs_joints" not in levels + + +def test_an_unknown_embodiment_stops_before_the_array_rules(tmp_path) -> None: + path = tmp_path / "eva.zarr" + _write_eva(path) + store = __import__("zarr").open_group(str(path), mode="a") + store.attrs["embodiment"] = "sharpa_bimanual" + + report = validate_episode(path) + + levels = _levels(report) + assert levels["attrs.embodiment"] == ERROR + assert levels["embodiment"] == ERROR + assert "left.obs_ee_pose" not in levels + + +def test_a_morphology_block_must_agree_with_the_embodiment_name(tmp_path) -> None: + path = tmp_path / "eva.zarr" + _write_eva(path) + store = __import__("zarr").open_group(str(path), mode="a") + store.attrs["morphology"] = { + "platform": "human_body", + "end_effector": "mano_hand", + } + + report = validate_episode(path) + + assert _levels(report)["attrs.morphology"] == ERROR + assert "disagrees" in report.text() + + +def test_a_morphology_block_selects_the_end_effector(tmp_path) -> None: + path = tmp_path / "eva.zarr" + _write_eva(path) + store = __import__("zarr").open_group(str(path), mode="a") + store.attrs["morphology"] = { + "platform": "eva_x5", + "end_effector": {"left": "eva_parallel_jaw", "right": "eva_parallel_jaw"}, + "vendor": "rl2", + } + + report = validate_episode(path) + + assert report.ok, report.text() + assert "eva_parallel_jaw" in report.text(verbose=True) + + +def test_a_directory_that_is_not_an_episode_reports_one_error(tmp_path) -> None: + report = validate_episode(tmp_path / "absent.zarr") + + assert not report.ok + assert _levels(report)["episode"] == ERROR + + +def test_the_cli_exit_code_follows_the_findings(tmp_path, capsys) -> None: + path = tmp_path / "eva.zarr" + _write_eva(path) + + assert main([str(path)]) == 0 + assert main([str(path), "--camera-coverage"]) == 0 + assert main([str(tmp_path / "absent.zarr")]) == 1 + assert "cannot open as a zarr group" in capsys.readouterr().out + + +def test_the_cli_can_report_json(tmp_path, capsys) -> None: + path = tmp_path / "eva.zarr" + _write_eva(path) + + main([str(path), "--json"]) + + payload = __import__("json").loads(capsys.readouterr().out) + assert payload[0]["ok"] is True + assert payload[0]["path"] == str(path) + assert payload[0]["requirements"] == dict.fromkeys(waivable_rules(), True) + + +@pytest.mark.parametrize("section", ["attributes", "checks", "arrays"]) +def test_every_schema_rule_declares_a_usable_requirement(section) -> None: + schema = load_schema() + rules = schema[section] + entries = rules.values() if isinstance(rules, dict) else rules + assert entries + for rule in entries: + assert isinstance(rule.get("required", False), bool) + assert rule["severity"] in ("integrity", "adoption", "coverage") + + +def test_waivable_rules_own_both_cli_flags_and_explain_why() -> None: + rules = waivable_rules() + assert set(rules) == { + "schema_version", + "data_status", + "calibration_block", + "camera_coverage", + "annotation_coverage", + "rgb_timestamps", + } + + help_text = " ".join(_build_parser().format_help().split()) + for name, rule in rules.items(): + flag = name.replace("_", "-") + assert f"--{flag}" in help_text + assert f"--no-{flag}" in help_text + assert " ".join(rule["why"].split()) in help_text + + args = _build_parser().parse_args( + ["episode.zarr", "--no-camera-coverage", "--annotation-coverage"] + ) + assert args.require_camera_coverage is False + assert args.require_annotation_coverage is True + assert args.require_schema_version is True + + +def test_integrity_rules_cannot_be_waived(tmp_path) -> None: + with pytest.raises(SchemaError, match="unknown waivable rule 'pose_degeneracy'"): + validate_episode( + tmp_path / "episode.zarr", + requirements={"pose_degeneracy": False}, + ) + + +def test_rule_levels_follow_the_invocation_without_weakening_integrity( + tmp_path, +) -> None: + schema = load_schema() + requirements = dict.fromkeys(waivable_rules(schema), True) + report = Report(path=tmp_path / "episode.zarr", requirements=requirements) + + calibration = schema["attributes"]["calibration"] + assert _level(calibration, report, "calibration") == ERROR + report.requirements["calibration_block"] = False + assert _level(calibration, report, "calibration") == WARNING + + integrity = schema["checks"][0] + assert integrity["name"] == "calibration_present" + assert _level(integrity, report, integrity["name"]) == ERROR + + +def test_the_retired_strict_flag_is_rejected() -> None: + with pytest.raises(SystemExit): + _build_parser().parse_args(["episode.zarr", "--strict"]) + + +# -------------------------------------------------------------------------- +# Degeneracy, timestamps and annotations +# -------------------------------------------------------------------------- + + +def _finding(report, check: str) -> str: + return next(f.message for f in report.findings if f.check == check) + + +def test_a_pose_track_that_never_moves_is_an_error(tmp_path) -> None: + numeric = _eva_numeric() + numeric["left.obs_ee_pose"] = np.tile(_poses()[0], (LENGTH, 1)) + _write_eva(tmp_path / "still.zarr", numeric=numeric) + + report = validate_episode(tmp_path / "still.zarr") + + assert _levels(report)["pose_degeneracy"] == ERROR + assert "constant across all 4 frames" in _finding(report, "pose_degeneracy") + + +def test_identity_rotations_beyond_the_limit_are_an_error(tmp_path) -> None: + numeric = _eva_numeric() + identity = numeric["left.obs_ee_pose"].copy() + identity[:, 3:7] = [1.0, 0.0, 0.0, 0.0] + numeric["left.obs_ee_pose"] = identity + _write_eva(tmp_path / "identity.zarr", numeric=numeric) + + report = validate_episode(tmp_path / "identity.zarr") + + assert _levels(report)["pose_degeneracy"] == ERROR + assert "identity rotation on 100% of frames" in _finding(report, "pose_degeneracy") + + +def test_an_identity_extrinsic_is_an_error(tmp_path) -> None: + path = tmp_path / "identity_rig.zarr" + ZarrWriter.create_and_write( + episode_path=path, + numeric_data=_eva_numeric(), + embodiment="eva_bimanual", + chunk_timesteps=LENGTH, + annotations=[("fold the towel", 0, LENGTH)], + intrinsics={"front_1": K}, + extrinsics={"left": np.eye(4), "right": np.eye(4)}, + ) + + report = validate_episode(path) + + assert _levels(report)["calibration_degeneracy"] == ERROR + assert "left arm base" in _finding(report, "calibration_degeneracy") + + +def test_a_synthesized_camera_matrix_is_an_error(tmp_path) -> None: + width = height = 480 + synthetic = np.array( + [ + [width, 0.0, width / 2, 0.0], + [0.0, width, height / 2, 0.0], + [0.0, 0.0, 1.0, 0.0], + ] + ) + _write_eva( + tmp_path / "synthetic.zarr", + calibration={ + "reference_frame": "camera:front_1", + "cameras": { + "front_1": {"K": synthetic.tolist(), "resolution": [width, height]} + }, + }, + ) + + report = validate_episode(tmp_path / "synthetic.zarr") + + assert _levels(report)["intrinsics_signature"] == ERROR + assert "synthesized centred pinhole" in _finding(report, "intrinsics_signature") + + +def test_a_rectified_aria_camera_matrix_stays_out_of_the_net(tmp_path) -> None: + # ``fx`` is 266.5 at width 640, so the full synthetic signature is false. + aria = np.array( + [[266.5, 0.0, 320.0, 0.0], [0.0, 266.5, 240.0, 0.0], [0.0, 0.0, 1.0, 0.0]] + ) + _write_eva( + tmp_path / "aria.zarr", + calibration={ + "reference_frame": "camera:front_1", + "cameras": {"front_1": {"K": aria.tolist(), "resolution": [640, 480]}}, + }, + ) + + report = validate_episode(tmp_path / "aria.zarr") + + assert _levels(report)["intrinsics_signature"] == OK + + +def test_a_stalled_clock_is_an_error(tmp_path) -> None: + numeric = _eva_numeric() + numeric["obs_rgb_timestamps_ns"] = np.array( + [1_000, 2_000, 2_000, 1_500], dtype=np.int64 + ) + _write_eva(tmp_path / "clock.zarr", numeric=numeric) + + report = validate_episode(tmp_path / "clock.zarr") + + assert _levels(report)["timestamps"] == ERROR + assert "does not increase on 2 of 3 steps" in _finding(report, "timestamps") + + +def test_a_float64_quantized_clock_is_an_error(tmp_path) -> None: + numeric = _eva_numeric() + numeric["obs_rgb_timestamps_ns"] = ( + np.int64(1_700_000_000_000_000_000) + np.arange(LENGTH, dtype=np.int64) * 256 + ) + _write_eva(tmp_path / "quantized.zarr", numeric=numeric) + + report = validate_episode(tmp_path / "quantized.zarr") + + assert "quantized to 256 ns" in _finding(report, "timestamps") + + +def test_a_timestamp_waiver_does_not_waive_a_wrong_dtype(tmp_path) -> None: + numeric = _eva_numeric() + numeric["obs_rgb_timestamps_ns"] = np.arange(LENGTH, dtype=np.float64) + _write_eva(tmp_path / "float_clock.zarr", numeric=numeric) + + report = validate_episode( + tmp_path / "float_clock.zarr", + requirements={"rgb_timestamps": False}, + ) + + assert _levels(report)["obs_rgb_timestamps_ns"] == ERROR + assert "expected int" in _finding(report, "obs_rgb_timestamps_ns") + + +def test_a_second_time_base_is_an_error(tmp_path) -> None: + numeric = _eva_numeric() + numeric["relative_timestamp_s"] = np.linspace(0.0, 0.1, LENGTH)[:, None] + _write_eva(tmp_path / "two_clocks.zarr", numeric=numeric) + + report = validate_episode(tmp_path / "two_clocks.zarr") + + assert _levels(report)["timestamps"] == ERROR + assert "one clock per episode" in _finding(report, "timestamps") + + +def test_annotation_coverage_below_the_minimum_is_reported(tmp_path) -> None: + _write_eva( + tmp_path / "thin.zarr", annotations=[("fold the towel", 0, LENGTH - 2)] + ) + + required = validate_episode(tmp_path / "thin.zarr") + waived = validate_episode( + tmp_path / "thin.zarr", + requirements={"annotation_coverage": False}, + ) + + assert _levels(required)["annotation_coverage"] == ERROR + assert _levels(waived)["annotation_coverage"] == WARNING + assert "cover 50% of the episode" in _finding(waived, "annotation_coverage") + + +def test_an_annotation_span_past_the_episode_is_reported(tmp_path) -> None: + _write_eva(tmp_path / "over.zarr", annotations=[("fold", 0, LENGTH + 5)]) + + report = validate_episode(tmp_path / "over.zarr") + + assert "outside [0, 4)" in _finding(report, "annotation_coverage") + + +def test_delimiter_encoded_metadata_is_an_error(tmp_path) -> None: + _write_eva( + tmp_path / "skill.zarr", + annotations=[("pick up the cup | Skill: pick", 0, LENGTH)], + ) + + report = validate_episode(tmp_path / "skill.zarr") + + assert _levels(report)["annotation_text"] == ERROR + assert "Skill: pick" in _finding(report, "annotation_text") + + +def test_delimiter_encoded_metadata_in_the_task_description_is_an_error( + tmp_path, +) -> None: + _write_eva(tmp_path / "desc.zarr", task_description="fold | Skill: fold") + + report = validate_episode(tmp_path / "desc.zarr") + + assert _levels(report)["annotation_text"] == ERROR + assert "task_description" in _finding(report, "annotation_text") diff --git a/egomimic/rldb/zarr/validate.py b/egomimic/rldb/zarr/validate.py new file mode 100644 index 000000000..82c8ee302 --- /dev/null +++ b/egomimic/rldb/zarr/validate.py @@ -0,0 +1,899 @@ +"""Validate a Zarr episode against ``schema/episode_v3.yaml``. + +Run it with:: + + python -m egomimic.rldb.zarr.validate [RULE FLAGS] + +The schema declares attributes, arrays, conditions, thresholds, and severity +classes. Integrity failures are always errors. Adoption and coverage failures +are errors by default, and each such rule supplies a matched ``--`` / +``--no-`` CLI pair. Validation of a present attribute or array is always +an error when its declared type or shape is wrong. +""" + +from __future__ import annotations + +import argparse +import functools +import json +import re +import sys +from collections.abc import Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np +import yaml +import zarr + +from egomimic.rldb.embodiment.embodiment import ( + Embodiment, + ResolvedEmbodiment, + canonical_embodiment_name, +) +from egomimic.rldb.embodiment.registry import load_embodiment_platforms +from egomimic.rldb.zarr.calibration import ( + IMAGE_KEY_PREFIX, + CalibrationError, + read_calibration, + uncalibrated_cameras, +) + +SCHEMA_DIR = Path(__file__).parent / "schema" +SCHEMA_FILE = SCHEMA_DIR / "episode_v3.yaml" + +ERROR = "error" +WARNING = "warning" +OK = "ok" + +#: Accepted values for a schema rule's ``severity`` class. +_SEVERITY_VALUES = ("integrity", "adoption", "coverage") +_INTEGRITY = "integrity" +_RULE_NAME = re.compile(r"^[a-z][a-z0-9_]*$") + +_TYPE_NAMES = { + "str": str, + "int": int, + "bool": bool, + "mapping": Mapping, +} + + +class SchemaError(ValueError): + """Raised when ``episode_v3.yaml`` contains an unsupported declaration.""" + + +@dataclass(frozen=True) +class Finding: + """One validation result emitted for an episode. + + Attributes: + level: ``ok``, ``warning``, or ``error``. + check: The attribute, array key, or named predicate being reported. + message: What the rule found, in one line. + """ + + level: str + check: str + message: str + + def __str__(self) -> str: + return f"{self.level.upper():<7} {self.check}: {self.message}" + + +@dataclass +class Report: + """Validation findings for one episode. + + Attributes: + path: The episode directory. + findings: Findings in validation order. A schema rule may emit zero, + one, or multiple findings after key expansion. + requirements: Whether each waivable rule is required. A false value + means that the invocation reports that rule's failures as warnings. + """ + + path: Path + findings: list[Finding] = field(default_factory=list) + requirements: dict[str, bool] = field(default_factory=dict) + + def add(self, level: str, check: str, message: str) -> None: + self.findings.append(Finding(level, check, message)) + + @property + def errors(self) -> list[Finding]: + return [f for f in self.findings if f.level == ERROR] + + @property + def warnings(self) -> list[Finding]: + return [f for f in self.findings if f.level == WARNING] + + @property + def ok(self) -> bool: + return not self.errors + + def summary(self) -> str: + return ( + f"{len(self.findings)} checks, {len(self.errors)} errors, " + f"{len(self.warnings)} warnings" + ) + + def text(self, verbose: bool = False) -> str: + """Format this report for terminal output. + + Args: + verbose: If true, list the rules that passed as well. + """ + lines = [str(self.path)] + for finding in self.findings: + if finding.level != OK or verbose: + lines.append(f" {finding}") + lines.append(f" {self.summary()}") + return "\n".join(lines) + + def to_jsonable(self) -> dict[str, Any]: + return { + "path": str(self.path), + "requirements": self.requirements, + "ok": self.ok, + "findings": [ + {"level": f.level, "check": f.check, "message": f.message} + for f in self.findings + ], + } + + +@functools.lru_cache(maxsize=1) +def load_schema() -> dict: + """Load the schema and validate its requirement declarations. + + Returns: + The parsed schema. + + Raises: + SchemaError: If a rule has an invalid requirement or severity class. + """ + with SCHEMA_FILE.open("r") as f: + schema = yaml.safe_load(f) or {} + waivable_names = set() + for default_name, rule in _schema_rules(schema): + required = rule.get("required", False) + if not isinstance(required, bool): + raise SchemaError( + f"{default_name!r}: `required` must be true or false, got {required!r}" + ) + severity = rule.get("severity") + if severity not in _SEVERITY_VALUES: + raise SchemaError( + f"{default_name!r}: `severity` must be one of " + f"{list(_SEVERITY_VALUES)}, got {severity!r}" + ) + if severity == _INTEGRITY: + continue + if not required: + raise SchemaError( + f"{default_name!r}: a waivable rule must be required by default" + ) + name = _rule_name(default_name, rule) + if not _RULE_NAME.fullmatch(name): + raise SchemaError( + f"{default_name!r}: rule name {name!r} must use snake_case" + ) + if name in waivable_names: + raise SchemaError(f"duplicate waivable rule name {name!r}") + waivable_names.add(name) + if not isinstance(rule.get("why"), str) or not rule["why"].strip(): + raise SchemaError( + f"{default_name!r}: a waivable rule needs a non-empty `why`" + ) + return schema + + +def _schema_rules(schema: Mapping): + """Yield ``(default_name, rule)`` pairs in schema order.""" + yield from schema.get("attributes", {}).items() + yield from ((rule.get("name"), rule) for rule in schema.get("checks", [])) + yield from ((rule.get("key"), rule) for rule in schema.get("arrays", [])) + + +def _rule_name(default_name: str, rule: Mapping) -> str: + """Return the stable name used by requirement overrides and CLI flags.""" + return rule.get("name", default_name) + + +def waivable_rules(schema: Mapping | None = None) -> dict[str, dict]: + """Return the schema rules whose failure severity callers may lower. + + The schema owns the rule names and explanations so the validator cannot add + a waivable rule without also exposing a documented CLI flag for it. + """ + schema = load_schema() if schema is None else schema + return { + _rule_name(default_name, rule): rule + for default_name, rule in _schema_rules(schema) + if rule.get("severity") != _INTEGRITY + } + + +def _requirements( + schema: Mapping, overrides: Mapping[str, bool] | None +) -> dict[str, bool]: + """Build one explicit error-or-warning decision per waivable rule.""" + requirements = dict.fromkeys(waivable_rules(schema), True) + for name, required in (overrides or {}).items(): + if name not in requirements: + raise SchemaError( + f"unknown waivable rule {name!r}; expected one of " + f"{sorted(requirements)}" + ) + if not isinstance(required, bool): + raise SchemaError( + f"requirement for {name!r} must be a bool, got {required!r}" + ) + requirements[name] = required + return requirements + + +def _level(rule: Mapping, report: Report, default_name: str) -> str | None: + """Return a missing or failed requirement's invocation-specific level.""" + if not rule.get("required", False): + return None + if rule["severity"] == _INTEGRITY: + return ERROR + name = _rule_name(default_name, rule) + return ERROR if report.requirements[name] else WARNING + + +# --------------------------------------------------------------------------- +# Attribute rules +# --------------------------------------------------------------------------- + + +def _check_type(value, type_name: str | None) -> str | None: + if type_name is None: + return None + expected = _TYPE_NAMES.get(type_name) + if expected is None: + raise SchemaError(f"unknown type {type_name!r} in the schema") + # A bool is an int in Python; the schema means them separately. + if expected is int and isinstance(value, bool): + return "expected an int, got a bool" + if not isinstance(value, expected): + return f"expected {type_name}, got {type(value).__name__}" + return None + + +def _check_attribute(name: str, rule: dict, attrs: Mapping, report, context) -> None: + if name not in attrs: + level = _level(rule, report, name) + if level is not None: + report.add(level, f"attrs.{name}", "missing") + return + + value = attrs[name] + problem = _check_type(value, rule.get("type")) + if problem is None and "min" in rule and value < rule["min"]: + problem = f"must be at least {rule['min']}, got {value}" + if problem is None and "choices" in rule and value not in rule["choices"]: + problem = f"expected one of {rule['choices']}, got {value!r}" + if problem is None and rule.get("check"): + problem = _ATTRIBUTE_CHECKS[rule["check"]](value, attrs, context) + + if problem is None: + report.add(OK, f"attrs.{name}", _describe(value)) + else: + report.add(ERROR, f"attrs.{name}", problem) + + +def _describe(value) -> str: + if isinstance(value, Mapping): + return f"{len(value)} entries" + text = str(value) + return text if len(text) <= 60 else text[:57] + "..." + + +def _check_embodiment_name(value, attrs, context) -> str | None: + if canonical_embodiment_name(value) not in load_embodiment_platforms(): + return ( + f"{value!r} is not owned by any platform in registry/platforms.yaml; " + f"known: {sorted(load_embodiment_platforms())}" + ) + return None + + +def _check_morphology(value, attrs, context) -> str | None: + try: + resolved = Embodiment.resolve(value) + except (TypeError, ValueError) as exc: + return str(exc) + named = context.get("named_platform") + if named is not None and resolved.platform.name != named.name: + return ( + f"platform {resolved.platform.name!r} disagrees with the platform " + f"{named.name!r} that embodiment {attrs.get('embodiment')!r} selects" + ) + return None + + +def _check_calibration(value, attrs, context) -> str | None: + try: + read_calibration(attrs) + except CalibrationError as exc: + return str(exc) + return None + + +def _check_schema_version(value, attrs, context) -> str | None: + known = load_schema().get("known_versions") or [] + if known and value not in known: + return f"unknown schema_version {value!r}; known: {known}" + return None + + +_ATTRIBUTE_CHECKS = { + "embodiment_name": _check_embodiment_name, + "morphology": _check_morphology, + "calibration": _check_calibration, + "schema_version": _check_schema_version, +} + + +# --------------------------------------------------------------------------- +# Named checks +# --------------------------------------------------------------------------- + + +def _check_calibration_present(rule, context, report) -> None: + calibration = context.get("calibration") + level = _level(rule, report, "calibration_present") + if calibration is None or not calibration.intrinsics(): + if level is not None: + report.add( + level, + "calibration_present", + "the episode states no camera matrix; calibration measures the " + "rig that recorded it and cannot be recovered later", + ) + return + report.add( + OK, + "calibration_present", + f"{len(calibration.intrinsics())} calibrated camera(s), " + f"reference_frame={calibration.reference_frame}", + ) + + +def _check_camera_coverage(rule, context, report) -> None: + missing = uncalibrated_cameras(context["array_keys"], context.get("calibration")) + level = _level(rule, report, "camera_coverage") + if missing: + if level is not None: + report.add( + level, + "camera_coverage", + f"no camera matrix for image stream(s) {missing}", + ) + return + report.add(OK, "camera_coverage", "every image stream has a camera matrix") + + +_IDENTITY = np.eye(4) + +#: Quaternion layout is `[qw, qx, qy, qz]`; both signs are the same rotation. +_IDENTITY_QUATERNIONS = ( + np.array([1.0, 0.0, 0.0, 0.0]), + np.array([-1.0, 0.0, 0.0, 0.0]), +) + + +def _read(array, total_frames: int | None) -> np.ndarray: + """Read the available prefix, capped at the authoritative frame count.""" + end = array.shape[0] if total_frames is None else min(total_frames, array.shape[0]) + return np.asarray(array[:end]) + + +def _report_problems(rule, report, check: str, problems, passed: str) -> None: + """Record a named check's pass or its requirement-level failure.""" + if not problems: + report.add(OK, check, passed) + return + level = _level(rule, report, check) + if level is not None: + report.add(level, check, "; ".join(problems)) + + +def _check_pose_degeneracy(rule, context, report) -> None: + suffixes = tuple(rule.get("suffixes") or ()) + threshold = float(rule.get("identity_rotation_fraction", 0.01)) + total_frames = context.get("total_frames") + problems = [] + checked = 0 + for key, array in context["arrays"].items(): + if not key.endswith(suffixes) or len(array.shape) != 2: + continue + poses = _read(array, total_frames) + if poses.shape[0] < 2 or poses.shape[1] != 7: + continue + checked += 1 + if len(np.unique(poses, axis=0)) == 1: + problems.append(f"{key} is constant across all {poses.shape[0]} frames") + rotations = poses[:, 3:7] + identity = np.zeros(len(rotations), dtype=bool) + for quaternion in _IDENTITY_QUATERNIONS: + identity |= np.all(rotations == quaternion, axis=1) + fraction = float(identity.mean()) + if fraction > threshold: + problems.append( + f"{key} holds an exact identity rotation on " + f"{fraction:.0%} of frames (limit {threshold:.0%})" + ) + _report_problems( + rule, report, "pose_degeneracy", problems, f"{checked} pose track(s) move" + ) + + +def _check_calibration_degeneracy(rule, context, report) -> None: + calibration = context.get("calibration") + if calibration is None: + return + problems = [] + for name, camera in calibration.cameras.items(): + # Identity is implicit for the reference camera; only an explicitly + # stored pose for another camera is subject to this branch. + if ( + camera.ref_T_cam is not None + and name != calibration.reference_camera + and np.array_equal(camera.ref_T_cam, _IDENTITY) + ): + problems.append(f"cameras[{name!r}].ref_T_cam is exactly the identity") + for side in calibration.arm_bases: + base_T_cam = calibration.base_T_cam(side) + if base_T_cam is not None and np.allclose(base_T_cam, _IDENTITY, atol=0.0): + problems.append( + f"the camera sits exactly at the {side} arm base " + "(base_T_cam is the identity)" + ) + _report_problems( + rule, report, "calibration_degeneracy", problems, "no identity extrinsic" + ) + + +def _resolution(camera, context) -> tuple[int, int] | None: + """Return ``(width, height)`` from calibration or image feature metadata.""" + if camera.resolution is not None: + return camera.resolution + feature = context["features"].get(f"{IMAGE_KEY_PREFIX}{camera.name}") or {} + shape = feature.get("shape") or [] + if len(shape) >= 2: + return int(shape[1]), int(shape[0]) + return None + + +def _check_intrinsics_signature(rule, context, report) -> None: + calibration = context.get("calibration") + if calibration is None: + return + problems = [] + checked = 0 + for name, camera in calibration.cameras.items(): + resolution = _resolution(camera, context) + if camera.K is None or resolution is None: + continue + checked += 1 + width, height = resolution + fx, fy = camera.K[0, 0], camera.K[1, 1] + cx, cy = camera.K[0, 2], camera.K[1, 2] + if fx == fy == width and cx == width / 2 and cy == height / 2: + problems.append( + f"cameras[{name!r}].K is a synthesized centred pinhole " + f"(fx = fy = {fx:g} = W, principal point at the image centre)" + ) + _report_problems( + rule, + report, + "intrinsics_signature", + problems, + f"{checked} camera matri(ces) carry a measured focal length", + ) + + +def _check_timestamps(rule, context, report) -> None: + problems = [] + for substring in rule.get("banned_key_substrings") or (): + stored = [k for k in context["arrays"] if substring in k] + if stored: + problems.append( + f"{stored} store a second time base; keep one clock per episode " + "and derive a relative base on read" + ) + + key = rule.get("key", "obs_rgb_timestamps_ns") + array = context["arrays"].get(key) + if array is None: + _report_problems(rule, report, "timestamps", problems, "no clock stored") + return + + stamps = _read(array, context.get("total_frames")) + if stamps.ndim != 1 or stamps.shape[0] < 2: + _report_problems(rule, report, "timestamps", problems, "one stamp") + return + + steps = np.diff(stamps.astype(np.int64)) + stalled = int(np.count_nonzero(steps <= 0)) + if stalled: + problems.append(f"{key} does not increase on {stalled} of {len(steps)} steps") + + quantum = int(np.gcd.reduce(np.abs(steps))) if steps.size else 0 + if quantum >= 64 and quantum & (quantum - 1) == 0: + problems.append( + f"{key} is quantized to {quantum} ns, the signature of float64 " + "seconds converted to nanoseconds; compute the clock in integers" + ) + + _report_problems( + rule, + report, + "timestamps", + problems, + f"{len(stamps)} stamps increase strictly", + ) + + +def _annotations(context, key: str) -> list[dict]: + array = context["arrays"].get(key) + if array is None: + return [] + out = [] + for entry in np.asarray(array[:]): + if isinstance(entry, (bytes, bytearray, memoryview)): + entry = bytes(entry).decode("utf-8") + if isinstance(entry, str): + try: + entry = json.loads(entry) + except json.JSONDecodeError: + continue + if isinstance(entry, Mapping): + out.append(dict(entry)) + return out + + +def _check_annotation_coverage(rule, context, report) -> None: + key = rule.get("key", "annotations") + total_frames = context.get("total_frames") or 0 + annotations = _annotations(context, key) + problems = [] + covered = np.zeros(max(total_frames, 0), dtype=bool) + for annotation in annotations: + start = int(annotation.get("start_idx", -1)) + end = int(annotation.get("end_idx", -1)) + if start < 0 or end > total_frames or end <= start: + problems.append( + f"span [{start}, {end}) is outside [0, {total_frames}) or empty" + ) + continue + covered[start:end] = True + fraction = float(covered.mean()) if covered.size else 0.0 + minimum = float(rule.get("minimum", 0.9)) + if fraction < minimum: + problems.append( + f"{len(annotations)} annotation(s) cover {fraction:.0%} of the " + f"episode (minimum {minimum:.0%}); trim the tail or annotate it" + ) + _report_problems( + rule, + report, + "annotation_coverage", + problems, + f"{len(annotations)} annotation(s) cover {fraction:.0%}", + ) + + +def _check_annotation_text(rule, context, report) -> None: + delimiters = tuple(rule.get("banned_delimiters") or ()) + texts = [ + (f"annotation {i}", a.get("text", "")) + for i, a in enumerate(_annotations(context, rule.get("key", "annotations"))) + ] + texts.append(("task_description", context["attrs"].get("task_description") or "")) + problems = [] + for where, text in texts: + for delimiter in delimiters: + if delimiter in str(text): + problems.append( + f"{where} encodes metadata after {delimiter!r}: {text!r}" + ) + break + _report_problems( + rule, report, "annotation_text", problems, "no delimiter-encoded metadata" + ) + + +_NAMED_CHECKS = { + "calibration_present": _check_calibration_present, + "camera_coverage": _check_camera_coverage, + "pose_degeneracy": _check_pose_degeneracy, + "calibration_degeneracy": _check_calibration_degeneracy, + "intrinsics_signature": _check_intrinsics_signature, + "timestamps": _check_timestamps, + "annotation_coverage": _check_annotation_coverage, + "annotation_text": _check_annotation_text, +} + + +# --------------------------------------------------------------------------- +# Array rules +# --------------------------------------------------------------------------- + + +def _condition_holds(when: Mapping, resolved: ResolvedEmbodiment, side) -> bool: + for name, expected in when.items(): + if name == "platform_kind": + actual = resolved.platform.kind + elif name == "has_arm_chain": + actual = resolved.platform.arm_dof is not None + elif name == "has_aux_chain": + actual = resolved.platform.aux is not None + elif name == "end_effector_class": + end_effector = resolved.end_effectors.get(side) if side else None + actual = None if end_effector is None else end_effector.ee_class + else: + raise SchemaError(f"unknown `when` condition {name!r} in the schema") + if actual != expected: + return False + return True + + +def _dimension(token, context, side) -> int | None: + """Resolve a shape token; return ``None`` for a wildcard or absent spec.""" + if isinstance(token, int): + return token + if token == "*": + return None + resolved: ResolvedEmbodiment = context["resolved"] + if token == "T": + return context["total_frames"] + if token == "arm_dof": + return resolved.platform.arm_dof + if token == "aux_dof": + return None if resolved.platform.aux is None else resolved.platform.aux.dof + if token == "ee_dof": + end_effector = resolved.end_effectors.get(side) if side else None + return None if end_effector is None else end_effector.dof + if token == "kp3": + end_effector = resolved.end_effectors.get(side) if side else None + if end_effector is None: + return None + return 3 * end_effector.keypoints.n_slots + raise SchemaError(f"unknown shape dimension {token!r} in the schema") + + +def _dtype_kind(dtype) -> str: + kind = getattr(dtype, "kind", "") + if kind in "iu": + return "int" + if kind == "f": + return "float" + if kind in "OSV": + return "object" + return kind or "any" + + +def _check_array(key: str, rule: dict, arrays: Mapping, report, context, side) -> None: + if key not in arrays: + level = _level(rule, report, key) + if level is not None: + report.add(level, key, "missing") + return + + array = arrays[key] + shape = tuple(int(n) for n in array.shape) + expected = rule.get("shape") + problems = [] + if expected is not None: + if len(shape) != len(expected): + problems.append( + f"expected {len(expected)} dimension(s), got shape {shape}" + ) + else: + for axis, token in enumerate(expected): + want = _dimension(token, context, side) + if want is None: + continue + # ``total_frames`` is the authoritative length and a stored + # array may carry a padded tail, so axis 0 is a lower bound. + if axis == 0: + if shape[0] < want: + problems.append( + f"holds {shape[0]} frames for total_frames {want}" + ) + elif shape[axis] != want: + problems.append( + f"axis {axis} is {shape[axis]}, expected {want} ({token})" + ) + + want_dtype = rule.get("dtype") + if want_dtype and want_dtype != "any": + actual = _dtype_kind(array.dtype) + if actual != want_dtype: + problems.append(f"dtype is {array.dtype}, expected {want_dtype}") + + if problems: + report.add(ERROR, key, "; ".join(problems)) + else: + report.add(OK, key, f"shape {shape} {array.dtype}") + + +def _expand_key(template: str, arrays: Mapping, sides) -> list[tuple[str, str | None]]: + """Expand a schema key into the concrete array keys it selects. + + ``{side}`` expands over the resolved side candidates. ``*`` selects stored + keys only, so a wildcard rule never creates a missing-key finding. + """ + if "{side}" in template: + candidates = [(template.format(side=side), side) for side in sides] + else: + candidates = [(template, None)] + out = [] + for key, side in candidates: + if "*" not in key: + out.append((key, side)) + continue + prefix, _, suffix = key.partition("*") + out.extend( + (name, side) + for name in arrays + if name.startswith(prefix) and name.endswith(suffix) + ) + return out + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + + +def validate_episode( + path: str | Path, *, requirements: Mapping[str, bool] | None = None +) -> Report: + """Validate one Zarr episode against ``schema/episode_v3.yaml``. + + Args: + path: The episode ``.zarr`` directory. + requirements: Per-rule severity decisions. ``True`` reports a failure + as an error; ``False`` waives it to a warning. Omitted rules remain + required. Integrity rules cannot be overridden. + + Returns: + The findings emitted while validating the episode. + + Raises: + SchemaError: If the schema or a requirement override is invalid. + """ + path = Path(path) + schema = load_schema() + report = Report(path=path, requirements=_requirements(schema, requirements)) + try: + store = zarr.open_group(str(path), mode="r") + except Exception as exc: + report.add(ERROR, "episode", f"cannot open as a zarr group: {exc}") + return report + + attrs = dict(store.attrs) + arrays = {name: store[name] for name in store.array_keys()} + + context: dict[str, Any] = { + "array_keys": list(arrays), + "arrays": arrays, + "attrs": attrs, + "features": attrs.get("features") or {}, + "total_frames": attrs.get("total_frames"), + } + embodiment = attrs.get("embodiment") + if embodiment: + context["named_platform"] = load_embodiment_platforms().get( + canonical_embodiment_name(embodiment) + ) + try: + context["calibration"] = read_calibration(attrs) + except CalibrationError: + context["calibration"] = None + + for name, rule in schema.get("attributes", {}).items(): + _check_attribute(name, rule, attrs, report, context) + + for rule in schema.get("checks", []): + check = _NAMED_CHECKS.get(rule.get("name")) + if check is None: + raise SchemaError(f"unknown check {rule.get('name')!r} in the schema") + check(rule, context, report) + + resolved = _resolve(attrs, report) + if resolved is None: + return report + # Record the resolved platform and end-effectors used by subsequent array + # conditions and dimensions. + report.add(OK, "embodiment", resolved.describe()) + context["resolved"] = resolved + + for rule in schema.get("arrays", []): + for key, side in _expand_key(rule["key"], arrays, resolved.sides): + if side is not None and side not in resolved.end_effectors: + continue + when = rule.get("when") or {} + if when and not _condition_holds(when, resolved, side): + continue + _check_array(key, rule, arrays, report, context, side) + + return report + + +def _resolve(attrs: Mapping, report: Report) -> ResolvedEmbodiment | None: + """Resolve morphology first, then fall back to the embodiment name.""" + for spec in (attrs.get("morphology"), attrs.get("embodiment")): + if not spec: + continue + try: + return Embodiment.resolve(spec) + except (TypeError, ValueError): + continue + report.add( + ERROR, + "embodiment", + "cannot resolve the episode's embodiment, so no array rule can run", + ) + return None + + +def _build_parser() -> argparse.ArgumentParser: + """Build the CLI parser from the schema's waivable rules.""" + parser = argparse.ArgumentParser( + prog="python -m egomimic.rldb.zarr.validate", + description=__doc__.splitlines()[0], + ) + parser.add_argument("paths", nargs="+", type=Path, help="episode .zarr paths") + for name, rule in waivable_rules().items(): + flag = name.replace("_", "-") + # argparse performs %-interpolation on help text. Schema rationales are + # ordinary prose and may contain percentages such as the 90% coverage + # threshold, so escape them before handing the text to argparse. + help_text = rule["why"].replace("%", "%%") + parser.add_argument( + f"--{flag}", + dest=f"require_{name}", + action=argparse.BooleanOptionalAction, + default=True, + help=f"[{rule['severity']}] {help_text} (default: required)", + ) + parser.add_argument( + "-v", + "--verbose", + action="store_true", + help="list the rules that passed as well", + ) + parser.add_argument( + "--json", action="store_true", help="print one JSON report per episode" + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + """Validate CLI paths and return an exit status. + + Returns: + ``0`` if every report has no errors, otherwise ``1``. + """ + parser = _build_parser() + args = parser.parse_args(argv) + requirements = {name: getattr(args, f"require_{name}") for name in waivable_rules()} + reports = [validate_episode(p, requirements=requirements) for p in args.paths] + if args.json: + print(json.dumps([r.to_jsonable() for r in reports], indent=2)) + else: + for report in reports: + print(report.text(verbose=args.verbose)) + return 0 if all(r.ok for r in reports) else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index a4d5d0a1b..c9c4642b1 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -43,11 +43,15 @@ # from action_chunk_transforms import Transform from egomimic.rldb.filters import DatasetFilter +from egomimic.rldb.zarr.action_chunk_transforms import base_T_cam_pose_key +from egomimic.rldb.zarr.calibration import Calibration, read_calibration +from egomimic.rldb.zarr.episode_attrs import data_status, is_complete from egomimic.utils.aws.aws_data_utils import load_env from egomimic.utils.aws.aws_sql import ( create_default_engine, episode_table_to_df, ) +from egomimic.utils.pose_utils import _matrix_to_xyzwxyz if TYPE_CHECKING: # Annotation-only import — avoids a runtime circular import with @@ -231,6 +235,16 @@ def _load_zarr_datasets(self, search_path: Path, valid_folder_names: set[str]): key_map=self.key_map, transform_list=self.transform_list, ) + # Resolver-based loading accepts only the exact ``complete`` + # status; structural samples and unknown statuses are skipped. + if not is_complete(ds_obj.metadata): + logger.warning( + "Skipping %s: data_status is %r, not 'complete'", + p, + ds_obj.data_status, + ) + skipped.append(p.name) + continue datasets[name] = ds_obj except Exception as e: logger.error(f"Failed to load dataset at {p}: {e}") @@ -1558,6 +1572,41 @@ def init_episode(self): self.keys_dict = {k: (0, None) for k in self.episode_reader._collect_keys()} self._image_keys = self._detect_image_keys() self._json_keys = self._detect_json_keys() + self._extrinsic_poses = self._build_extrinsic_poses() + + @property + def data_status(self) -> str: + """Return the episode status, including the legacy default.""" + return self.episode_reader.data_status + + @property + def calibration(self) -> Calibration | None: + """Return parsed current or legacy episode calibration.""" + return self.episode_reader.calibration + + def _build_extrinsic_poses(self) -> dict[str, np.ndarray]: + """Build transform-input poses for the calibration's default camera. + + For each arm with enough calibration to compose ``base_T_cam``, convert + the matrix to the pose layout consumed by the transform pipeline. + Missing calibration leaves the sample key absent so the configured + transform may supply its compatibility fallback. + + Returns: + A mapping from ``_base_T_cam_pose`` to a 7-value + ``[xyz, qw, qx, qy, qz]`` pose. + """ + calibration = self.calibration + if calibration is None: + return {} + poses = {} + for side in ("left", "right"): + base_T_cam = calibration.base_T_cam(side) + if base_T_cam is not None: + poses[base_T_cam_pose_key(side)] = _matrix_to_xyzwxyz( + base_T_cam[None, :] + )[0] + return poses @property def intrinsics(self) -> np.ndarray | dict[str, np.ndarray] | None: @@ -1725,6 +1774,11 @@ def _next(reason: str, key: str = "") -> int: if retry: continue + # Inject episode-derived poses before transforms run. Transform + # fallbacks use ``setdefault`` and therefore preserve these values. + for key, pose in self._extrinsic_poses.items(): + data[key] = pose.copy() + if self.transform: for transform in self.transform or []: data = transform.transform(data) @@ -1734,29 +1788,16 @@ def _next(reason: str, key: str = "") -> int: data[k] = torch.from_numpy(v).to(torch.float32) data["embodiment"] = get_embodiment_id(self.embodiment) - # Per-episode camera intrinsics travel with the batch so a single - # data-driven embodiment can project without a hardcoded class const. - # Single-camera -> (3,4) K matrix; no/multi-camera -> NaN sentinel, - # which _intrinsics_from_batch treats as "fall back to cls.INTRINSICS". - K = self.intrinsics - if isinstance(K, dict): - # Multi-camera rig: project against the front camera (viz/projection - # target the front image). Prefer a "front"-keyed entry, else the - # first available camera. - K = next( - (v for k, v in K.items() if "front" in str(k).lower()), - next(iter(K.values()), None) if K else None, - ) - if K is not None: - K = np.asarray(K, dtype=np.float32) - if K.shape == (3, 3): - # Normalize a 3x3 K to the canonical 3x4 (zeros last column); - # some contributors store 3x3 (e.g. microagi). - K = np.concatenate([K, np.zeros((3, 1), dtype=np.float32)], axis=1) - if K.shape != (3, 4): # unexpected -> sentinel (viz falls back to const) - K = np.full((3, 4), np.nan, dtype=np.float32) - else: + # ``Calibration.K()`` applies the default-camera selection and the + # parser has already normalized a 3x3 K to 3x4. A missing matrix is + # represented by NaNs, which ``_intrinsics_from_batch`` interprets + # as a request for the embodiment-class fallback. + calibration = self.calibration + K = None if calibration is None else calibration.K() + if K is None: K = np.full((3, 4), np.nan, dtype=np.float32) + else: + K = np.asarray(K, dtype=np.float32) data["intrinsics"] = torch.from_numpy(np.ascontiguousarray(K)) ep_name = Path(self.episode_path).name data["episode_hash"] = ep_name[:-5] if ep_name.endswith(".zarr") else ep_name @@ -1829,6 +1870,7 @@ class ZarrEpisode: "_store", "metadata", "keys", + "_calibration", ) def __init__(self, path: str | Path): @@ -1841,6 +1883,21 @@ def __init__(self, path: str | Path): self._store = zarr.open_group(str(self._path), mode="r") self.metadata = dict(self._store.attrs) self.keys = self.metadata["features"] + self._calibration = read_calibration(self.metadata) + + @property + def data_status(self) -> str: + """Return the status, defaulting a missing or falsey value to ``complete``.""" + return data_status(self.metadata) + + @property + def calibration(self) -> Calibration | None: + """Return calibration normalized by :func:`read_calibration`. + + The result comes from the current block or a lifted legacy attribute + pair. It is ``None`` when neither representation is present. + """ + return self._calibration @property def intrinsics(self) -> np.ndarray | dict[str, np.ndarray] | None: @@ -1848,16 +1905,12 @@ def intrinsics(self) -> np.ndarray | dict[str, np.ndarray] | None: Camera intrinsics persisted in zarr metadata, deserialized to ndarray(s). Returns: - - np.ndarray for a single-camera episode, - - dict[str, np.ndarray] for multi-camera episodes, - - None if no intrinsics were written. + - dict[str, np.ndarray] keyed by camera name, + - None if the episode carries no calibrated camera. """ - raw = self.metadata.get("intrinsics") - if raw is None: + if self._calibration is None: return None - if isinstance(raw, dict): - return {k: np.asarray(v) for k, v in raw.items()} - return np.asarray(raw) + return self._calibration.intrinsics() or None def read( self, keys_with_ranges: dict[str, tuple[int, int | None]] diff --git a/egomimic/rldb/zarr/zarr_writer.py b/egomimic/rldb/zarr/zarr_writer.py index 39b375f60..58b80f079 100644 --- a/egomimic/rldb/zarr/zarr_writer.py +++ b/egomimic/rldb/zarr/zarr_writer.py @@ -6,6 +6,7 @@ """ import json +import logging from pathlib import Path from typing import Any, Literal @@ -14,6 +15,19 @@ import zarr from zarr.core.dtype import VariableLengthBytes +from egomimic.rldb.zarr.calibration import ( + Calibration, + lift_legacy_calibration, + parse_calibration, + uncalibrated_cameras, +) +from egomimic.rldb.zarr.episode_attrs import ( + DATA_STATUS_COMPLETE, + DATA_STATUS_VALUES, +) + +logger = logging.getLogger(__name__) + def _intrinsics_to_jsonable( intrinsics: np.ndarray | list | dict, @@ -282,6 +296,8 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> bool: padding[:] = last_jpeg self._store[key][self._total_frames : padded] = padding + self._writer._check_camera_coverage(self._image_info) + # Write language annotations if self._writer.annotations is not None: self._writer._write_annotations(self._store, self._writer.annotations) @@ -309,6 +325,9 @@ def __init__( chunk_timesteps: int = 100, intrinsics: dict | None = None, extrinsics: dict | None = None, + calibration: Calibration | dict | None = None, + data_status: str = DATA_STATUS_COMPLETE, + require_camera_coverage: bool = False, verbose: bool = False, ): """Configure a writer for one Zarr v3 episode. @@ -323,9 +342,16 @@ def __init__( chunk_timesteps: Number of frames in each numeric-array chunk. intrinsics: A mapping from camera keys to 3×4 camera matrices. This low-level constructor permits ``None``. - extrinsics: ``None`` or a mapping from keys to 4×4 ``ref_T_cam`` - matrices. Robot episodes use arm names as keys. See - ``docs/CONVENTIONS.md``. + extrinsics: The legacy ``{side: base_T_cam}`` mapping, or ``None``. + See ``docs/CONVENTIONS.md``. + calibration: A ``Calibration`` or its attribute-block mapping. The + writer stores it under ``calibration``. + data_status: Initial episode status. The staging path and dataset + resolvers accept only ``complete`` and exclude + ``structural_sample``. + require_camera_coverage: If true, reject an image stream with no + matching camera matrix; otherwise log a warning. Coverage is + named explicitly because this option changes no other check. verbose: If true, print progress information during writes. """ self.episode_path = Path(episode_path) @@ -339,6 +365,16 @@ def __init__( self.chunk_timesteps = chunk_timesteps self.intrinsics = intrinsics self.extrinsics = extrinsics + self.calibration = ( + None if calibration is None else parse_calibration(calibration) + ) + if data_status not in DATA_STATUS_VALUES: + raise ValueError( + f"data_status must be one of {list(DATA_STATUS_VALUES)}, got " + f"{data_status!r}" + ) + self.data_status = data_status + self.require_camera_coverage = require_camera_coverage self.verbose = verbose # Track image shapes for metadata self._features: dict[str, dict[str, Any]] = {} @@ -419,6 +455,10 @@ def write( store, key, enc_arr, img_shape, padded_frames ) + self._check_camera_coverage( + list(image_data) + list(pre_encoded_image_data) + ) + # Write language annotations if provided if self.annotations is not None: self._write_annotations(store, self.annotations) @@ -691,6 +731,37 @@ def _write_annotations( "format": "annotation_v1", } + def _check_camera_coverage(self, image_keys) -> None: + """Compare stored image streams with the effective camera calibration. + + The current ``calibration`` block is authoritative when present; + otherwise the legacy ``intrinsics`` and ``extrinsics`` attributes are + lifted. Each ``images.`` key must match a camera carrying ``K``. + + Args: + image_keys: The episode's image array keys. + + Raises: + ValueError: If camera coverage is required and a stream carries no + ``K``. + """ + calibration = self.calibration or lift_legacy_calibration( + self.intrinsics, self.extrinsics + ) + missing = uncalibrated_cameras(image_keys, calibration) + if not missing: + return + message = ( + f"{self.episode_path.name}: no camera matrix for image stream(s) " + f"{missing}. Every `images.` array needs a K in " + "`calibration.cameras`. See CONTRIBUTING_DATA.md §6.4." + ) + if self.require_camera_coverage: + raise ValueError(message) + logger.warning( + "%s Pass require_camera_coverage=True to make this an error.", message + ) + def _build_metadata( self, metadata_override: dict[str, Any] | None = None ) -> dict[str, Any]: @@ -705,6 +776,7 @@ def _build_metadata( """ metadata = { "embodiment": self.embodiment, + "data_status": self.data_status, "total_frames": self.total_frames, "fps": self.fps, "task_name": self.task_name, @@ -712,22 +784,22 @@ def _build_metadata( "features": self._features, } + if self.calibration is not None: + metadata["calibration"] = self.calibration.to_jsonable() + if self.intrinsics is not None: metadata["intrinsics"] = _intrinsics_to_jsonable(self.intrinsics) if self.extrinsics is not None: metadata["extrinsics"] = _intrinsics_to_jsonable(self.extrinsics) - # Apply overrides — but NEVER let them clobber the validated camera - # metadata. intrinsics/extrinsics are validated in create_and_write; a - # converter's metadata_override that happens to carry a stale or empty - # "intrinsics"/"extrinsics" key must not silently overwrite the validated - # values (this was the Mecka clobber bug → empty intrinsics in zarr.json). + # Preserve the current and legacy calibration representations supplied + # to the writer; metadata overrides may add fields but not replace them. if metadata_override: override = { k: v for k, v in metadata_override.items() - if k not in ("intrinsics", "extrinsics") + if k not in ("calibration", "intrinsics", "extrinsics") } metadata.update(override) @@ -746,8 +818,11 @@ def create_and_write( annotations: list[tuple[str, int, int]] | None = None, chunk_timesteps: int = 100, *, - intrinsics: dict, + intrinsics: dict | None = None, extrinsics: dict | None = None, + calibration: Calibration | dict | None = None, + data_status: str = DATA_STATUS_COMPLETE, + require_camera_coverage: bool = False, metadata_override: dict[str, Any] | None = None, ) -> Path: """Validate episode metadata and write one Zarr v3 episode. @@ -767,12 +842,19 @@ def create_and_write( annotations: ``(text, start_idx, end_idx)`` annotation tuples. chunk_timesteps: Number of frames in each numeric-array chunk. intrinsics: A non-empty mapping from camera keys to 3×4 camera - matrices. - extrinsics: ``None`` or a non-empty mapping from keys to 4×4 - ``ref_T_cam`` matrices. Robot episodes use arm names as keys. - See ``docs/CONVENTIONS.md``. + matrices. Optional only when ``calibration`` supplies them. + extrinsics: The legacy non-empty ``{side: base_T_cam}`` mapping, or + ``None``. See ``docs/CONVENTIONS.md``. + calibration: A ``Calibration`` or its attribute-block mapping. For + each omitted legacy argument, the writer derives an equivalent + ``intrinsics`` or ``extrinsics`` value for older readers. + data_status: ``complete`` or ``structural_sample``. The staging path + and dataset resolvers accept only ``complete`` episodes. + require_camera_coverage: If true, reject an image stream with no + matching camera matrix; otherwise log a warning. Coverage is + named explicitly because this option changes no other check. metadata_override: Additional episode metadata. This mapping cannot - replace ``intrinsics`` or ``extrinsics``. + replace ``calibration``, ``intrinsics``, or ``extrinsics``. Returns: The path of the created ``.zarr`` directory. @@ -783,6 +865,10 @@ def create_and_write( ValueError: If ``intrinsics`` is empty or contains a non-3×4 value. ValueError: If ``extrinsics`` is not ``None`` or a non-empty mapping of 4×4 matrices. + ValueError: If ``data_status`` is not a supported value. + ValueError: If camera coverage is required and an image stream + carries no camera matrix. + CalibrationError: If ``calibration`` is malformed. """ # Validate the embodiment up front (it is guaranteed to be consumed by # the reader via get_embodiment_id) so a bad/empty/typo'd value fails @@ -795,11 +881,21 @@ def create_and_write( f"{[m.name.lower() for m in EMBODIMENT]}, got {embodiment!r}. " "See CONTRIBUTING_DATA.md §9." ) + # Derive only omitted legacy fields. The current block remains the + # source used by readers that understand it. + if calibration is not None: + calibration = parse_calibration(calibration) + if intrinsics is None: + intrinsics = calibration.intrinsics() + if extrinsics is None: + extrinsics = calibration.extrinsics() or None + if not isinstance(intrinsics, dict) or not intrinsics: raise ValueError( "Camera intrinsics must be a non-empty {camera_key: 3x4 K matrix} " 'dict for a contributed zarr episode, e.g. {"front_1": K}. ' - "See CONTRIBUTING_DATA.md §6.4." + "Pass `intrinsics=` or a `calibration=` block that carries a K " + "for at least one camera. See CONTRIBUTING_DATA.md §6.4." ) for cam_key, K in intrinsics.items(): try: @@ -843,6 +939,9 @@ def create_and_write( chunk_timesteps=chunk_timesteps, intrinsics=intrinsics, extrinsics=extrinsics, + calibration=calibration, + data_status=data_status, + require_camera_coverage=require_camera_coverage, ) writer.write( diff --git a/egomimic/scripts/backfill_scripts/stage_processed_folder.py b/egomimic/scripts/backfill_scripts/stage_processed_folder.py index 67359f8e3..075e80d77 100644 --- a/egomimic/scripts/backfill_scripts/stage_processed_folder.py +++ b/egomimic/scripts/backfill_scripts/stage_processed_folder.py @@ -85,9 +85,13 @@ def find_zarr_prefixes(s3, base): def read_batch(prefixes, folder, cfg): - """Read the zarr.json for a batch of prefixes -> list[row dict]. FULLY - self-contained (only stdlib + boto3, imported inside) so it runs on any Ray - worker regardless of whether egomimic / ~/.egoverse_env are present there.""" + """Build staging rows from the episode metadata under each S3 prefix. + + The function imports its dependencies locally so a Ray worker needs neither + the ``egomimic`` package nor ``~/.egoverse_env``. It omits every episode + whose status is not exactly ``complete``; a missing or falsey status + defaults to ``complete`` for compatibility with older episodes. + """ import json as _json import re as _re from datetime import datetime as _dt @@ -121,6 +125,10 @@ def read_batch(prefixes, folder, cfg): ).replace(tzinfo=_tz.utc).isoformat() except ValueError: created_at = None + # Keep the literal aligned with ``DATA_STATUS_COMPLETE`` without + # importing the package on a bare Ray worker. + if (attrs.get("data_status") or "complete") != "complete": + continue feats = attrs.get("features") out.append({ "episode_hash": episode_hash, diff --git a/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py b/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py index 81551de95..be5032b6a 100644 --- a/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py +++ b/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py @@ -20,8 +20,9 @@ can't handle yet. Projection convention (human egocentric `human_bimanual`): the head IS the -camera. Per-episode intrinsics are read from `zarr.json` -(`grp.attrs["intrinsics"]["front_1"]`, a 3x4 K); poses in `*.obs_ee_pose` +camera. Per-episode intrinsics are read from `zarr.json` through +`read_calibration` (the `calibration` block, or the legacy `intrinsics` +attribute of an older episode); poses in `*.obs_ee_pose` (xyz + quat wxyz, SLAM-world) are transformed world → head frame using `obs_head_pose` (same xyz+wxyz layout) and projected with that K. There is no per-arm extrinsic. (This matches the validated `Human.viz` path; the old @@ -43,6 +44,7 @@ # `Embodiment.viz` method projects the poses and draws the overlay. # `projectaria_tools` is optional in the standalone visualization environment. # If an embodiment import fails, the browser shows an unavailable badge. +from egomimic.rldb.zarr.calibration import read_calibration from egomimic.utils.pose_utils import ee_pose_to_cam_frame from .images import ( @@ -283,49 +285,45 @@ def _badge(img_rgb, text: str): return img_rgb -def _intrinsics_from_zarr(grp): - """Per-episode 3x4 camera matrix K from zarr.json metadata - (`attrs["intrinsics"]["front_1"]`). A 3x3 K is normalized to 3x4 by - appending a zero column. Returns None if absent/unparseable.""" +def _calibration_from_zarr(grp): + """Read current or legacy calibration without breaking the inspector. + + Returns: + The normalized calibration, or ``None`` if metadata access fails, no + calibration is present, or its stored representation is malformed. + Invalid calibration disables the overlay instead of the episode view. + """ try: - intr = dict(grp.attrs).get("intrinsics") - if not intr: - return None - K = np.asarray(intr["front_1"], dtype=float) + return read_calibration(dict(grp.attrs)) except Exception: return None - if K.shape == (3, 4): - return K - if K.shape == (3, 3): - return np.hstack([K, np.zeros((3, 1))]) - return None + + +def _intrinsics_from_zarr(grp): + """Return the default camera's normalized 3×4 ``[K_3x3 | 0]`` matrix. + + Return ``None`` if the episode has no valid calibration or if the selected + camera has no matrix. + """ + calibration = _calibration_from_zarr(grp) + return None if calibration is None else calibration.K() def _extrinsics_from_zarr(grp): - """Read per-arm ``base_T_cam`` matrices from episode metadata. + """Compose per-arm poses of the calibration's default camera. Returns: - A dictionary with valid ``"left"`` and ``"right"`` 4×4 matrices. - The function returns ``None`` if no valid matrix is present. It ignores - missing, non-numeric, and incorrectly shaped values. + A dictionary mapping available ``"left"`` and ``"right"`` arms to + 4×4 ``base_T_cam`` matrices, or ``None`` if none can be composed. """ - try: - extr = dict(grp.attrs).get("extrinsics") - if not extr: - return None - except Exception: + calibration = _calibration_from_zarr(grp) + if calibration is None: return None out = {} for arm in ("left", "right"): - m = extr.get(arm) - if m is None: - continue - try: - M = np.asarray(m, dtype=float) - except Exception: - continue - if M.shape == (4, 4): - out[arm] = M + base_T_cam = calibration.base_T_cam(arm) + if base_T_cam is not None: + out[arm] = base_T_cam return out or None diff --git a/pyproject.toml b/pyproject.toml index 1b6471417..d46a86b10 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,9 +69,10 @@ dependencies = [ [tool.setuptools.packages.find] where = ["."] -# Include the registry YAML files in built distributions. +# Include the registry and schema YAML files in built distributions. [tool.setuptools.package-data] "egomimic.rldb.embodiment.registry" = ["*.yaml"] +"egomimic.rldb.zarr" = ["schema/*.yaml"] [tool.pytest.ini_options] testpaths = ["egomimic/rldb/embodiment", "egomimic/rldb/zarr"]