rpx_benchmark.data

189def features_for_task(task: TaskType | str) -> Features:
190    """Return the canonical :class:`datasets.Features` for ``task``.
191
192    Parameters
193    ----------
194    task : TaskType or str
195
196    Returns
197    -------
198    datasets.Features
199    """
200    if isinstance(task, str):
201        task = TaskType(task)
202    return RPX_FEATURES[task]

264@dataclass
265class RPXHFBridge:
266    """Iterable adapter: ``datasets.Dataset`` → :class:`Sample` batches.
267
268    Shaped like :class:`rpx_benchmark.loader.RPXDataset` (``__len__`` +
269    ``__iter__`` yielding ``list[Sample]``) so the runner can consume
270    either source interchangeably.
271
272    Parameters
273    ----------
274    hf_dataset : datasets.Dataset
275        A dataset with columns matching the canonical task schema.
276    task : TaskType
277    batch_size : int, default 1
278    """
279
280    hf_dataset: Any
281    task: TaskType
282    batch_size: int = 1
283
284    def __len__(self) -> int:
285        return int(len(self.hf_dataset))
286
287    def __iter__(self) -> Iterator[List[Sample]]:
288        batch: List[Sample] = []
289        for row in self.hf_dataset:
290            batch.append(row_to_sample(row, self.task))
291            if len(batch) == self.batch_size:
292                yield batch
293                batch = []
294        if batch:
295            yield batch

220def row_to_sample(row: Dict[str, Any], task: TaskType | str) -> Sample:
221    """Convert one HF ``Dataset`` row into a :class:`Sample`.
222
223    Parameters
224    ----------
225    row : dict
226        Single row fetched from ``hf_ds[i]``. Column layout must match
227        :func:`rpx_benchmark.data.features.features_for_task`.
228    task : TaskType or str
229
230    Returns
231    -------
232    Sample
233
234    Raises
235    ------
236    ManifestError
237        If required columns are missing or carry invalid shapes.
238    """
239    if isinstance(task, str):
240        task = TaskType(task)
241
242    metadata: Dict[str, Any] = {}
243    # Sibling RGB for pair tasks lives in ``rgb_b``; we preserve it so
244    # downstream adapters can see it without re-reading the dataset.
245    if row.get("rgb_b") is not None:
246        metadata["rgb_b"] = _decode_image(row["rgb_b"])
247
248    gt_builder = _GT_BUILDERS.get(task)
249    if gt_builder is None:
250        raise ManifestError(f"No HF bridge registered for task {task}.")
251    gt = gt_builder(row)
252
253    return Sample(
254        id=str(row["id"]),
255        rgb=_decode_image(row["rgb"]),
256        ground_truth=gt,
257        metadata=metadata or None,
258        phase=_enum_or_none(Phase, row.get("phase")),
259        difficulty=_enum_or_none(Difficulty, row.get("difficulty")),
260        camera_pose=_unflatten_pose(row.get("camera_pose")),
261    )

 46def load_hf(
 47    task: TaskType | str,
 48    split: Difficulty | str = Difficulty.HARD,
 49    *,
 50    repo_id: str = DEFAULT_REPO_ID,
 51    revision: str | None = DEFAULT_REVISION,
 52    cache_dir: str | None = None,
 53    streaming: bool = False,
 54    batch_size: int = 1,
 55    **load_dataset_kwargs: Any,
 56) -> RPXHFBridge:
 57    """Load an RPX task slice as an :class:`RPXHFBridge`.
 58
 59    Parameters
 60    ----------
 61    task : TaskType or str
 62        Task config name on the HF repo. Uses the enum ``.value`` if a
 63        :class:`TaskType` is passed.
 64    split : Difficulty or str, default ``"hard"``
 65        ESD split to fetch. Maps to ``datasets.load_dataset``'s
 66        ``split`` kwarg.
 67    repo_id : str
 68        Hugging Face dataset repo id. Defaults to
 69        ``IRVLUTD/RPX`` (overridable via ``RPX_HF_REPO``).
 70    revision : str, optional
 71        Git revision (tag, branch, or commit) to pin. Defaults to the
 72        consolidated audited release; override it with ``RPX_HF_REVISION``.
 73    cache_dir : str, optional
 74        Hugging Face cache directory. When None the library default is
 75        used (``~/.cache/huggingface``).
 76    streaming : bool, default False
 77        Pass through to ``datasets.load_dataset``. Enables row-level
 78        streaming from the Hub without materialising the full Arrow
 79        tables locally.
 80    batch_size : int, default 1
 81        Forwarded into the returned :class:`RPXHFBridge`.
 82    **load_dataset_kwargs
 83        Any extra kwargs are forwarded to ``datasets.load_dataset``.
 84
 85    Returns
 86    -------
 87    RPXHFBridge
 88        An iterable wrapper that yields ``list[Sample]`` batches.
 89
 90    Raises
 91    ------
 92    DownloadError
 93        If the underlying ``datasets.load_dataset`` call fails.
 94    """
 95    task_enum = TaskType(task) if isinstance(task, str) else task
 96    split_name = split.value if isinstance(split, Difficulty) else str(split)
 97
 98    log.info(
 99        "loading HF dataset repo=%s config=%s split=%s streaming=%s",
100        repo_id,
101        task_enum.value,
102        split_name,
103        streaming,
104    )
105    try:
106        hf_ds = _load_dataset(
107            repo_id,
108            name=task_enum.value,
109            split=split_name,
110            revision=revision,
111            cache_dir=cache_dir,
112            streaming=streaming,
113            **load_dataset_kwargs,
114        )
115    except Exception as e:
116        raise DownloadError(
117            f"datasets.load_dataset failed for {repo_id}/{task_enum.value} "
118            f"(split={split_name}): {e}",
119            hint=(
120                "Check the repo id and your HF_TOKEN for private datasets. "
121                "For offline reuse set HF_DATASETS_OFFLINE=1 after a first "
122                "online load populates the cache."
123            ),
124        ) from e
125
126    return RPXHFBridge(hf_dataset=hf_ds, task=task_enum, batch_size=batch_size)

130@dataclass
131class PhaseFeatures:
132    """Features and provenance for one (scene, phase) pair."""
133
134    scene_id: str  # full scene-dir basename
135    phase: int  # 0 / 1 / 2 (clutter / interaction / clean)
136    n_frames_total: int  # frames present in rgb/
137    n_frames_used: int  # frames where every required modality was present
138    features: Dict[str, float] = field(default_factory=dict)
139
140    def as_json(self) -> Dict[str, object]:
141        return asdict(self)

666def extract_phase_features(phase_dir: Path) -> PhaseFeatures:
667    """Compute all 18 ESD features for one (scene, phase) directory.
668
669    ``phase_dir`` is expected to be ``<data_root>/<scene_dir>/<phase_idx>``;
670    e.g. ``test_dataset_aggregated/scene1.GDC.cse/0``. The scene_id and
671    phase fields of the result are inferred from this path.
672
673    The per-frame loop is streaming: one frame's modalities are decoded,
674    fed into per-feature accumulators, and discarded. Peak memory is
675    independent of phase length.
676    """
677    phase_dir = Path(phase_dir)
678    scene_id, phase_idx = _scene_phase_from_dir(phase_dir)
679
680    rgb_dir = phase_dir / "rgb"
681    n_total = len(_stem_set(rgb_dir, ".png")) if rgb_dir.is_dir() else 0
682    stems = _aligned_frames(phase_dir)
683
684    if not stems:
685        log.warning("no aligned frames found for %s/%d", scene_id, phase_idx)
686        return PhaseFeatures(
687            scene_id=scene_id,
688            phase=phase_idx,
689            n_frames_total=n_total,
690            n_frames_used=0,
691            features=_empty_features(),
692        )
693
694    T = len(stems)
695
696    # Annotation-effort features need no per-frame I/O.
697    feats: Dict[str, float] = dict(_annotation_effort(phase_dir, stems))
698
699    # Streaming accumulators.
700    obj_counts = np.zeros(T, dtype=np.int32)
701    depth_invalid_frac = np.zeros(T, dtype=np.float64)
702    depth_inv_mask_frac = np.full(T, np.nan, dtype=np.float64)  # NaN → no mask pixels
703    depth_std_per_frame = np.full(T, np.nan, dtype=np.float64)  # NaN → all-invalid
704    depth_std_mask_per_frame = np.full(T, np.nan, dtype=np.float64)  # NaN → no valid depth in mask
705    spec_frac = np.zeros(T, dtype=np.float64)
706    dark_frac = np.zeros(T, dtype=np.float64)
707    rgb_blur_per_frame = np.zeros(T, dtype=np.float64)  # Laplacian variance
708    rgb_tex_per_frame = np.zeros(T, dtype=np.float64)  # gradient magnitude mean
709    bboxes_per_frame: List[Dict[int, Tuple[int, int, int, int]]] = []
710    # Per-frame dense bincount of mask values (length varies by frame max ID);
711    # zipped at the end into a (T, max_id) matrix for area_cv / vis_instability.
712    bincount_per_frame: List[np.ndarray] = []
713    positions = np.zeros((T, 3), dtype=np.float64)
714    quats = np.zeros((T, 4), dtype=np.float64)
715
716    for t, s in enumerate(stems):
717        depth_m = _load_depth_m(phase_dir / "depth" / f"{s}.png")
718        mask = _load_mask(phase_dir / "sam2/masks" / f"{s}.png")
719        rgb_luma = _load_rgb_luma(phase_dir / "rgb" / f"{s}.png")
720        pos, q = _load_pose(phase_dir / "cam_pose" / f"{s}.npz")
721
722        # ── Depth-quality features ────────────────────────────────────────
723        invalid = depth_m <= 0
724        depth_invalid_frac[t] = float(invalid.mean())
725
726        in_mask = mask > 0
727        if in_mask.any():
728            depth_inv_mask_frac[t] = float(invalid[in_mask].mean())
729
730        valid = depth_m[~invalid]
731        if valid.size > 0:
732            depth_std_per_frame[t] = float(valid.std())
733
734        # In-mask depth std: only valid-depth pixels that also fall inside any
735        # instance mask. Captures object-relevant depth noise (paper's
736        # depth_std is whole-image and dominated by background).
737        valid_in_mask = depth_m[(~invalid) & in_mask]
738        if valid_in_mask.size > 0:
739            depth_std_mask_per_frame[t] = float(valid_in_mask.std())
740
741        # ── Photometric-conflict features ─────────────────────────────────
742        spec_frac[t] = float(((rgb_luma > _SPECULAR_LUMA_MIN) & invalid).mean())
743        dark_frac[t] = float(((rgb_luma < _DARK_LUMA_MAX) & invalid).mean())
744
745        # ── Image quality (D435 RGB) — grounded in CEMS (Zhao et al. 2024) ──
746        rgb_blur_per_frame[t] = _laplacian_var(rgb_luma)
747        rgb_tex_per_frame[t] = _gradient_magnitude_mean(rgb_luma)
748
749        # ── Mask-derived features (one bincount per frame, reused 3×) ─────
750        flat = mask.ravel()
751        # Drop negative IDs (defensive — masks should be unsigned).
752        if flat.min() < 0:
753            flat = flat[flat >= 0]
754        counts = np.bincount(flat)  # counts[i] = pixel count for ID i
755        ids = np.flatnonzero(counts[1:]) + 1 if counts.size > 1 else np.empty(0, dtype=np.int64)
756        obj_counts[t] = ids.size
757        bincount_per_frame.append(counts)
758
759        bboxes: Dict[int, Tuple[int, int, int, int]] = {}
760        for oid in ids:
761            b = _bbox_xywh(mask == oid)
762            if b is not None:
763                bboxes[int(oid)] = b
764        bboxes_per_frame.append(bboxes)
765
766        # ── Pose ──────────────────────────────────────────────────────────
767        positions[t] = pos
768        quats[t] = q
769
770    # ── n_total_objects: prefer mask_to_object.json, else union of observed ──
771    obj_json = phase_dir / "sam2" / "mask_to_object.json"
772    if obj_json.is_file():
773        try:
774            n_total_objects = len(json.loads(obj_json.read_text()))
775        except (ValueError, OSError) as e:
776            log.warning("%s: cannot parse, falling back to observed IDs (%s)", obj_json, e)
777            n_total_objects = max(
778                (int(c.size - 1) for c in bincount_per_frame if c.size > 1),
779                default=0,
780            )
781    else:
782        all_ids = {
783            int(i)
784            for c in bincount_per_frame
785            for i in (np.flatnonzero(c[1:]) + 1 if c.size > 1 else ())
786        }
787        n_total_objects = len(all_ids)
788
789    # ── Scene-complexity features ────────────────────────────────────────
790    feats["obj_mean"] = float(obj_counts.mean())
791    feats["obj_std"] = float(obj_counts.std())
792    feats["obj_consist"] = (
793        float((obj_counts == n_total_objects).mean()) if n_total_objects > 0 else 0.0
794    )
795
796    # ── Occlusion ────────────────────────────────────────────────────────
797    feats.update(_occlusion_from_per_frame_bboxes(bboxes_per_frame))
798
799    # ── Depth aggregates (NaN means "skip frame for this aggregate") ─────
800    feats["depth_invalid"] = float(depth_invalid_frac.mean())
801    feats["depth_invalid_mask"] = (
802        float(np.nanmean(depth_inv_mask_frac)) if not np.isnan(depth_inv_mask_frac).all() else 0.0
803    )
804    feats["depth_std"] = (
805        float(np.nanmean(depth_std_per_frame)) if not np.isnan(depth_std_per_frame).all() else 0.0
806    )
807    feats["depth_std_mask"] = (
808        float(np.nanmean(depth_std_mask_per_frame))
809        if not np.isnan(depth_std_mask_per_frame).all()
810        else 0.0
811    )
812
813    # ── Photometric ──────────────────────────────────────────────────────
814    feats["specular"] = float(spec_frac.mean())
815    feats["dark"] = float(dark_frac.mean())
816
817    # ── Image quality (D435 RGB) ─────────────────────────────────────
818    feats["rgb_blur"] = float(rgb_blur_per_frame.mean())
819    feats["rgb_texture"] = float(rgb_tex_per_frame.mean())
820
821    # ── Compute max_id once for reuse by object-size and temporal-stability ─
822    max_id = max((int(c.size - 1) for c in bincount_per_frame if c.size > 1), default=0)
823
824    # ── Object size distribution ────────────────────────────────────
825    # Per-instance mean mask area fraction (over frames where visible),
826    # then mean/std across instances. Small objects are universally harder.
827    if max_id >= 1 and T > 0:
828        _frame_pixels = bincount_per_frame[0].sum() if bincount_per_frame else 1
829        _areas_for_size = np.zeros((T, max_id), dtype=np.int64)
830        for _t, _c in enumerate(bincount_per_frame):
831            _n = min(int(_c.size) - 1, max_id)
832            if _n > 0:
833                _areas_for_size[_t, :_n] = _c[1 : 1 + _n]
834        _present = _areas_for_size > 0
835        _per_inst_mean = []
836        for _j in range(max_id):
837            _vis = _areas_for_size[:, _j][_present[:, _j]]
838            if _vis.size > 0:
839                _per_inst_mean.append(float(_vis.mean()) / max(_frame_pixels, 1))
840        if _per_inst_mean:
841            feats["mask_area_mean"] = float(np.mean(_per_inst_mean))
842            feats["mask_area_std"] = float(np.std(_per_inst_mean))
843        else:
844            feats["mask_area_mean"] = 0.0
845            feats["mask_area_std"] = 0.0
846    else:
847        feats["mask_area_mean"] = 0.0
848        feats["mask_area_std"] = 0.0
849
850    # ── Temporal stability: build dense (T, max_id) area matrix ──────────
851    if max_id >= 1:
852        areas = np.zeros((T, max_id), dtype=np.int64)
853        for t, c in enumerate(bincount_per_frame):
854            n = min(int(c.size) - 1, max_id)
855            if n > 0:
856                areas[t, :n] = c[1 : 1 + n]
857        feats.update(_temporal_stability_from_areas(areas))
858    else:
859        feats["area_cv"] = 0.0
860        feats["vis_instability"] = 0.0
861
862    # ── Camera motion ────────────────────────────────────────────────────
863    feats.update(_camera_motion_from_arrays(positions, quats))
864
865    # ── Fisheye / stereo (independent I/O — uses fisheye/, not the
866    # rgb/depth/mask/pose modalities streamed above) ─────────────────────
867    feats.update(_fisheye_features(phase_dir))
868
869    # Stable column order matching FEATURE_NAMES.
870    # Missing fisheye features → NaN (not 0) so percentile normalization
871    # can exclude them instead of biasing toward low difficulty.
872    feats = {
873        k: float(feats.get(k, float("nan") if k in _NAN_WHEN_ABSENT else 0.0))
874        for k in FEATURE_NAMES
875    }
876
877    return PhaseFeatures(
878        scene_id=scene_id,
879        phase=phase_idx,
880        n_frames_total=n_total,
881        n_frames_used=T,
882        features=feats,
883    )

886def iter_phase_dirs(data_root: Path) -> Iterator[Path]:
887    """Yield every ``<scene>/<phase>`` directory under ``data_root``.
888
889    A phase directory is any immediate child of a scene directory whose
890    name is purely numeric and that contains an ``rgb/`` subdirectory.
891    """
892    data_root = Path(data_root)
893    if not data_root.is_dir():
894        raise DatasetError(
895            f"data_root is not a directory: {data_root}", hint="check the --data-root argument."
896        )
897    for scene_dir in sorted(p for p in data_root.iterdir() if p.is_dir()):
898        for child in sorted(scene_dir.iterdir()):
899            if child.is_dir() and child.name.isdigit() and (child / "rgb").is_dir():
900                yield child