Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion egomimic/eval/latent_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@
import logging

from egomimic.rldb.filters import DatasetFilter
from egomimic.rldb.zarr.zarr_dataset_multi import EvenStrideDataset, MultiDataset
from egomimic.rldb.multi import EvenStrideDataset, MultiDataset

logger = logging.getLogger(__name__)

Expand Down
2 changes: 1 addition & 1 deletion egomimic/hydra_configs/data/aria.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper

train_datasets:
human_bimanual:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
_target_: egomimic.rldb.multi.MultiDataset._from_resolver
resolver:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
folder_path: ${paths.dataset_dir}
Expand Down
2 changes: 1 addition & 1 deletion egomimic/hydra_configs/data/aria_keypoints.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper

train_datasets:
human_bimanual:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
_target_: egomimic.rldb.multi.MultiDataset._from_resolver
resolver:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
folder_path: ${paths.dataset_dir}
Expand Down
2 changes: 1 addition & 1 deletion egomimic/hydra_configs/data/aria_keypoints_wrist.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper

train_datasets:
human_bimanual:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
_target_: egomimic.rldb.multi.MultiDataset._from_resolver
resolver:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
folder_path: ${paths.dataset_dir}
Expand Down
2 changes: 1 addition & 1 deletion egomimic/hydra_configs/data/aria_pilot_lance.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper

train_datasets:
human_bimanual:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
_target_: egomimic.rldb.multi.MultiDataset._from_resolver
resolver:
_target_: egomimic.rldb.lance.dataset.LanceEpisodeResolver
lance_root: ${oc.env:EGOVERSE_LANCE_PILOT,/coc/flash7/scratch/lance_pilot}/lance
Expand Down
2 changes: 1 addition & 1 deletion egomimic/hydra_configs/data/aria_pilot_zarr.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper

train_datasets:
human_bimanual:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
_target_: egomimic.rldb.multi.MultiDataset._from_resolver
resolver:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver
folder_path: ${oc.env:EGOVERSE_LANCE_PILOT,/coc/flash7/scratch/lance_pilot}/zarr
Expand Down
2 changes: 1 addition & 1 deletion egomimic/hydra_configs/data/aria_qwen.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper

train_datasets:
human_bimanual:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
_target_: egomimic.rldb.multi.MultiDataset._from_resolver
resolver:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
folder_path: ${paths.dataset_dir}
Expand Down
4 changes: 2 additions & 2 deletions egomimic/hydra_configs/data/cotrain_pi_base.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper

train_datasets:
eva_bimanual:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
_target_: egomimic.rldb.multi.MultiDataset._from_resolver
resolver:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
folder_path: ${paths.dataset_dir}
Expand All @@ -33,7 +33,7 @@ train_datasets:
mode: train
valid_ratio: 0.05
human_bimanual:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
_target_: egomimic.rldb.multi.MultiDataset._from_resolver
resolver:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
folder_path: ${paths.dataset_dir}
Expand Down
2 changes: 1 addition & 1 deletion egomimic/hydra_configs/data/eva.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper

train_datasets:
eva_bimanual:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
_target_: egomimic.rldb.multi.MultiDataset._from_resolver
resolver:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
folder_path: ${paths.dataset_dir}
Expand Down
2 changes: 1 addition & 1 deletion egomimic/hydra_configs/data/human.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper
# ('aria_bimanual'). Until that migration lands, this dataset matches zero rows.
train_datasets:
human_bimanual: &hb
_target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
_target_: egomimic.rldb.multi.MultiDataset._from_resolver
resolver:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
folder_path: ${paths.dataset_dir}
Expand Down
2 changes: 1 addition & 1 deletion egomimic/hydra_configs/data/mecka.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper

train_datasets:
human_bimanual:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
_target_: egomimic.rldb.multi.MultiDataset._from_resolver
resolver:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
folder_path: ${paths.dataset_dir}
Expand Down
2 changes: 1 addition & 1 deletion egomimic/hydra_configs/data/mecka_all_6d.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper

train_datasets:
human_bimanual:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
_target_: egomimic.rldb.multi.MultiDataset._from_resolver
resolver:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
# Per-cluster zarr mirror (paths/default.yaml; $EGOVERSE_DATASET_DIR
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ valid_prefix: seen_op_valid

train_datasets:
human_bimanual:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
_target_: egomimic.rldb.multi.MultiDataset._from_resolver
resolver:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
folder_path: ${paths.dataset_dir}
Expand Down
2 changes: 1 addition & 1 deletion egomimic/hydra_configs/data/mecka_fold_freeform_6d.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper

train_datasets:
human_bimanual:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
_target_: egomimic.rldb.multi.MultiDataset._from_resolver
resolver:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
folder_path: ${paths.dataset_dir}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ held_out_operators:

train_datasets:
human_bimanual:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
_target_: egomimic.rldb.multi.MultiDataset._from_resolver
resolver:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
folder_path: ${paths.dataset_dir}
Expand Down
2 changes: 1 addition & 1 deletion egomimic/hydra_configs/data/scale.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper

train_datasets:
human_bimanual:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
_target_: egomimic.rldb.multi.MultiDataset._from_resolver
resolver:
_target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
folder_path: ${paths.dataset_dir}
Expand Down
2 changes: 1 addition & 1 deletion egomimic/pl_utils/pl_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@

import egomimic.utils.tensor_utils as TensorUtils
from egomimic.pl_utils.pl_data_utils import head_of_loader, video_loader_name
from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset
from egomimic.rldb.multi import MultiDataset


def _barrier() -> None:
Expand Down
2 changes: 1 addition & 1 deletion egomimic/rldb/embodiment/embodiment.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import numpy as np
import torch

from egomimic.rldb.zarr.action_chunk_transforms import Transform
from egomimic.rldb.transforms import Transform
from egomimic.utils.type_utils import _to_numpy
from egomimic.utils.viz_utils import (
_viz_annotations,
Expand Down
2 changes: 1 addition & 1 deletion egomimic/rldb/embodiment/eva.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
_strip_pi_keymap_mode,
)
from egomimic.rldb.embodiment.human import ARIA_INTRINSICS
from egomimic.rldb.zarr.action_chunk_transforms import (
from egomimic.rldb.transforms import (
ActionChunkCoordinateFrameTransform,
BatchQuaternionPoseToYPR,
CartesianRot6DToYPR,
Expand Down
2 changes: 1 addition & 1 deletion egomimic/rldb/embodiment/human.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
_reject_legacy_rotation,
_strip_pi_keymap_mode,
)
from egomimic.rldb.zarr.action_chunk_transforms import (
from egomimic.rldb.transforms import (
ActionChunkCoordinateFrameTransform,
BatchQuaternionPoseToYPR,
CartesianRot6DToYPR,
Expand Down
16 changes: 16 additions & 0 deletions egomimic/rldb/filters.py
Original file line number Diff line number Diff line change
Expand Up @@ -145,3 +145,19 @@ def matches(self, row: Mapping[str, Any]) -> bool:
if row.get("episode_hash") not in self.completed_episode_hashes:
return False
return super().matches(row)


class PinError(ValueError):
"""A pinned episode_hash cannot be used: missing, deleted, no processed
path, or its embodiment does not match the dataset it was pinned under."""


def _ensure_dataset_filter(filters: DatasetFilter | None) -> DatasetFilter:
if filters is None:
return DatasetFilter()
if isinstance(filters, DatasetFilter):
return filters
raise TypeError(
"filters must be a DatasetFilter or None; plain dict filters are no longer "
"supported."
)
143 changes: 143 additions & 0 deletions egomimic/rldb/lance/annotations.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,143 @@
"""Review of annotation change files before `commit.apply_annotations` commits them.

A change file is a Lance or parquet table with an `op` column ("add" or "delete") and the
ANNOTATIONS fields. An add row may leave `annotation_id` null, since it is computed from
the row's content; a delete row needs only `annotation_id`.
"""

from collections import defaultdict
from dataclasses import dataclass, field
from pathlib import Path

import lance
import pyarrow as pa
import pyarrow.parquet as pq

from egomimic.rldb.lance.commit import _exists, annotations_uri

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit (naming): _exists is now used outside commit.py, and the new aria_to_lance_ray.py repeats the same (uri / "_versions").exists() check at lines 83 and 115. Make it public (e.g. commit.table_exists) and use it in both places.

from egomimic.rldb.lance.schema import ANNOTATION_TYPES, ANNOTATIONS, annotation_id

CONTENT = ("recording_id", "annotation_vendor", "annotation_type", "t_start_ns", "t_end_ns", "text") # fmt: skip


@dataclass
class Review:
adds: pa.Table
deletes: list[str]
errors: list[str] = field(default_factory=list)
unchanged: list[str] = field(default_factory=list)


def read_changes(path: str | Path) -> pa.Table:
path = str(path).rstrip("/")
return (
lance.dataset(path).to_table()
if path.endswith(".lance")
else pq.read_table(path)
)


def committed(root: Path) -> pa.Table:
uri = annotations_uri(root)
if not _exists(uri):
return ANNOTATIONS.empty_table()
return lance.dataset(str(uri)).to_table()


def _task_bounds(root: Path) -> dict[str, list[tuple[int, int]]]:
cols = ["recording_id", "t_start_ns", "t_end_ns"]
out = defaultdict(list)
for t in (
lance.dataset(str(Path(root) / "tasks.lance"))
.to_table(columns=cols)
.to_pylist()
):
out[t["recording_id"]].append((t["t_start_ns"], t["t_end_ns"]))
return out


def _add_errors(r: dict, tasks: dict) -> list[str]:
errs = [f"{k} is missing" for k in CONTENT if r.get(k) in (None, "")]
if errs:
return errs
if r["annotation_type"] not in ANNOTATION_TYPES:
errs.append(
f"annotation_type {r['annotation_type']!r} not in {ANNOTATION_TYPES}"
)
t0, t1 = r["t_start_ns"], r["t_end_ns"]
if t0 >= t1:
errs.append(f"t_start_ns {t0} is not before t_end_ns {t1}")
if r["recording_id"] not in tasks:
errs.append(f"recording {r['recording_id']} has no tasks")
elif not any(s <= t0 and t1 <= e for s, e in tasks[r["recording_id"]]):

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium (correctness): This containment check rejects any annotation that covers a recording's last frame. Annotations are half-open, so such a span ends at ts[-1] + period. But aria_to_lance writes the task's t_end_ns as int(ts[-1]) (inclusive), and lance_pilot/commit.py copies that into tasks.lance. migrate-zarr produces exactly these end-of-clip spans, so review flags them and commit refuses the whole file. Please pick one convention: either make the task/alias end half-open in aria_to_lance (ts[-1] + period), or allow t1 <= e + period here. Also make tests/fixtures/lance_tables.py match what aria_to_lance actually writes.

errs.append(f"[{t0}, {t1}) is not inside any task of its recording")
q = r.get("quality_score")
if q is not None and not 0 <= q <= 5:
errs.append(f"quality_score {q} is outside 0-5")
return errs


def review(root: Path, changes: pa.Table) -> Review:
"""Validate `changes` against the committed tables; nothing is written."""
tasks = _task_bounds(root)
have = set(committed(root).column("annotation_id").to_pylist())
adds, deletes, errors, unchanged, seen = [], [], [], [], set()
for i, r in enumerate(changes.to_pylist()):
op, given = r.get("op"), r.get("annotation_id")
if op == "delete":
if not given:
errors.append(f"row {i}: delete without annotation_id")
elif given not in have:
errors.append(f"row {i}: delete of {given}, which is not committed")
elif given in seen:
errors.append(f"row {i}: {given} appears twice")
else:
deletes.append(given)
seen.add(given)
continue
if op != "add":
errors.append(f"row {i}: op {op!r} is not add or delete")
continue
errs = _add_errors(r, tasks)
if errs:
errors += [f"row {i}: {e}" for e in errs]
continue
aid = annotation_id(*(r[k] for k in CONTENT))
if given and given != aid:
errors.append(
f"row {i}: annotation_id {given} does not match its content ({aid})"
)
elif aid in seen:
errors.append(f"row {i}: {aid} appears twice")
elif aid in have:
unchanged.append(aid)
else:
adds.append(
{**{k: r.get(k) for k in ANNOTATIONS.names}, "annotation_id": aid}
)
seen.add(aid)
return Review(
pa.Table.from_pylist(adds, schema=ANNOTATIONS), deletes, errors, unchanged
)


def diff(root: Path, rv: Review) -> str:
"""Per recording: rows added, deleted and already committed, with their texts."""
old = {r["annotation_id"]: r for r in committed(root).to_pylist()}
per = defaultdict(lambda: {"+": [], "-": [], "=": 0})
for r in rv.adds.to_pylist():
per[r["recording_id"]]["+"].append(r)
for aid in rv.deletes:
per[old[aid]["recording_id"]]["-"].append(old[aid])
for aid in rv.unchanged:
per[old[aid]["recording_id"]]["="] += 1
lines = []
for rid in sorted(per):
p = per[rid]
lines.append(f"{rid}: +{len(p['+'])} -{len(p['-'])} ={p['=']}")
for sign in "-+":
for r in sorted(p[sign], key=lambda r: r["t_start_ns"]):
lines.append(
f" {sign} [{r['t_start_ns']}, {r['t_end_ns']}) "
f"{r['annotation_vendor']}/{r['annotation_type']}: {r['text']}"
)
return "\n".join(lines)
Loading
Loading