rpx_benchmark.data.esd

  1"""Effort-Stratified Difficulty (ESD) feature extraction.
  2
  3Computes the 18 per-(scene, phase) features defined in
  4``paper-submission/overleaf/text/12_appendix.tex``
  5(``\\section{ESD Feature Definitions}``). The output of this module
  6is the raw feature table; weight calibration and tertile assignment
  7live downstream.
  8
  9On-disk layout expected
 10-----------------------
 11
 12::
 13
 14    <data_root>/
 15      <scene_dir>/                       # e.g. "scene1.GDC.cse" or just "scene1"
 16        0/  1/  2/                       # phase indices (clutter/interaction/clean)
 17          rgb/      00000.png ...        # uint8 RGB
 18          depth/    00000.png ...        # uint16 millimetres, 0 = invalid
 19          cam_pose/ 00000.npz            # T265 pose: position + orientation [x,y,z,w]
 20          sam2/
 21            masks/  00000.png            # 1-channel, pixel value > 0 = instance ID
 22            mask_to_object.json          # {"<id>": "<label>", ...} (optional)
 23            iter1_faulty.txt             # frame indices still faulty after iter i
 24            iter2_faulty.txt             # (one entry per line, basename or int)
 25            iter3_faulty.txt
 26            iter4_faulty.txt
 27
 28Frame indices align across modalities by basename
 29(``00000.png`` ↔ ``00000.npz``).
 30
 31Implementation notes
 32--------------------
 33
 34The per-phase loop is **streaming**: one frame's RGB / depth / mask /
 35pose is decoded, contributes to per-feature accumulators, then is
 36released. Peak memory per worker is ~tens of MB regardless of phase
 37length, so the CLI can default to ``os.cpu_count()`` workers.
 38
 39Mask geometry uses a single ``np.bincount`` per frame to derive
 40visible-IDs, per-instance areas, and the population for occlusion
 41bboxes — replacing the previous N-pass-over-mask design.
 42"""
 43
 44from __future__ import annotations
 45
 46import functools
 47import json
 48import math
 49import re
 50from dataclasses import asdict, dataclass, field
 51from pathlib import Path
 52from typing import Dict, Iterator, List, Optional, Tuple
 53
 54import numpy as np
 55from PIL import Image
 56
 57from ..exceptions import DatasetError
 58from ..logging_utils import get_logger
 59
 60log = get_logger(__name__)
 61
 62# Optional cv2 fast-path for PNG decode. Falls back to PIL if cv2 isn't
 63# installed; tests run on PIL so behaviour stays identical.
 64try:
 65    import cv2  # type: ignore
 66
 67    _HAS_CV2 = True
 68except ImportError:  # pragma: no cover
 69    _HAS_CV2 = False
 70
 71
 72# Names of the 31 features, in the order they appear in the appendix.
 73# Tests should diff against this list to catch silent additions / renames.
 74FEATURE_NAMES: Tuple[str, ...] = (
 75    # Annotation effort
 76    "iter_mean",
 77    "iter_max",
 78    # Scene complexity
 79    "obj_mean",
 80    "obj_std",
 81    "obj_consist",
 82    # Occlusion
 83    "occ_mean",
 84    "occ_p90",
 85    "occ_heavy",
 86    # Depth quality
 87    "depth_invalid",
 88    "depth_invalid_mask",
 89    "depth_std",
 90    "depth_std_mask",
 91    # Photometric--depth conflict (D435 RGB + invalid depth)
 92    "specular",
 93    "dark",
 94    # Image quality (D435 RGB) — grounded in CEMS (Zhao et al. 2024)
 95    "rgb_blur",
 96    "rgb_texture",
 97    # Object size distribution
 98    "mask_area_mean",
 99    "mask_area_std",
100    # Temporal annotation stability
101    "area_cv",
102    "area_drop",
103    "vis_instability",
104    # Camera motion (MAD-filtered)
105    "trans_mean",
106    "trans_p90",
107    "rot_mean",
108    "rot_p90",
109    "jerk",
110    # Fisheye / stereo (T265 fisheye pairs) — NaN if no fisheye dir present
111    "fisheye_dark",
112    "fisheye_bright",
113    "fisheye_sharpness",
114    "fisheye_corr",
115    "fisheye_texture",
116)
117
118# Threshold used by ``occ-heavy`` (appendix §"Occlusion").
119_OCC_HEAVY_THRESHOLD: float = 0.3
120# Photometric thresholds used by ``specular`` / ``dark`` / ``fisheye_{dark,bright}``.
121_SPECULAR_LUMA_MIN: int = 230
122_DARK_LUMA_MAX: int = 30
123# Threshold used by ``area_drop`` (old paper §3.4 eq 15).
124# Relative area collapse from frame t-1 to frame t exceeding this fraction counts
125# as a "drop event". 0.5 = 50% area lost between consecutive frames.
126_AREA_DROP_THRESHOLD: float = 0.5
127
128
129@dataclass
130class PhaseFeatures:
131    """Features and provenance for one (scene, phase) pair."""
132
133    scene_id: str  # full scene-dir basename
134    phase: int  # 0 / 1 / 2 (clutter / interaction / clean)
135    n_frames_total: int  # frames present in rgb/
136    n_frames_used: int  # frames where every required modality was present
137    features: Dict[str, float] = field(default_factory=dict)
138
139    def as_json(self) -> Dict[str, object]:
140        return asdict(self)
141
142
143# --------------------------------------------------------------------------- #
144# Frame discovery / loaders
145# --------------------------------------------------------------------------- #
146
147
148def _stem_set(d: Path, suffix: str) -> set[str]:
149    if not d.is_dir():
150        return set()
151    return {p.stem for p in d.iterdir() if p.suffix == suffix and not p.name.startswith(".")}
152
153
154def _aligned_frames(phase_dir: Path) -> List[str]:
155    """Frame stems that exist across rgb / depth / mask / pose, sorted.
156
157    Logs a warning summarising any per-modality misalignment so silent
158    data loss is visible without spamming the log per missing frame.
159    """
160    rgb = _stem_set(phase_dir / "rgb", ".png")
161    depth = _stem_set(phase_dir / "depth", ".png")
162    masks = _stem_set(phase_dir / "sam2/masks", ".png")
163    pose = _stem_set(phase_dir / "cam_pose", ".npz")
164    common = rgb & depth & masks & pose
165
166    union = rgb | depth | masks | pose
167    if union and len(common) < len(union):
168        missing = {
169            "rgb": len(union - rgb),
170            "depth": len(union - depth),
171            "mask": len(union - masks),
172            "pose": len(union - pose),
173        }
174        gaps = ", ".join(f"{k}:{v}" for k, v in missing.items() if v > 0)
175        log.warning(
176            "%s: dropped %d frames due to modality misalignment (missing in: %s)",
177            phase_dir,
178            len(union) - len(common),
179            gaps,
180        )
181    return sorted(common)
182
183
184@functools.lru_cache(maxsize=None)
185def _warn_multichannel_once(parent_dir: str, n_channels: int) -> None:
186    """Emit one multi-channel-PNG warning per parent directory.
187
188    A phase with hundreds of mistakenly-RGB-saved masks would otherwise spam
189    the log with one warning per frame; this collapses them to one line per
190    directory. ``lru_cache`` is the dedupe; tests call ``cache_clear()`` to
191    reset between cases.
192    """
193    log.warning(
194        "%s: multi-channel PNGs detected (%d ch); using first channel "
195        "(further warnings for this directory suppressed)",
196        parent_dir,
197        n_channels,
198    )
199
200
201def _decode_png_gray(path: Path) -> np.ndarray:
202    """Decode a single-channel PNG. Uses cv2 if available, PIL otherwise.
203
204    Multi-channel inputs (e.g. masks accidentally saved as RGB) are collapsed
205    to the first channel; one warning is emitted per containing directory
206    (see :func:`_warn_multichannel_once`).
207    """
208    if _HAS_CV2:
209        arr = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
210        if arr is None:
211            raise DatasetError(
212                f"failed to decode PNG: {path}", hint="file may be corrupt or zero-byte."
213            )
214        if arr.ndim == 3:
215            _warn_multichannel_once(str(path.parent), arr.shape[2])
216            arr = arr[..., 0]
217        return arr
218    img = Image.open(path)
219    arr = np.asarray(img)
220    if arr.ndim == 3:
221        _warn_multichannel_once(str(path.parent), arr.shape[2])
222        arr = arr[..., 0]
223    return arr
224
225
226def _load_depth_m(path: Path) -> np.ndarray:
227    """16-bit PNG millimetres → float32 metres; 0 stays 0 (sentinel)."""
228    arr = _decode_png_gray(path).astype(np.float32) / 1000.0
229    return arr
230
231
232def _load_mask(path: Path) -> np.ndarray:
233    """Single-channel PNG; pixel value > 0 = instance ID. Returned as int32."""
234    arr = _decode_png_gray(path)
235    return arr.astype(np.int32, copy=False)
236
237
238def _load_rgb_luma(path: Path) -> np.ndarray:
239    """RGB → ITU-R BT.601 luma (uint8) for photometric thresholds."""
240    if _HAS_CV2:
241        arr = cv2.imread(str(path), cv2.IMREAD_COLOR)  # BGR
242        if arr is None:
243            raise DatasetError(
244                f"failed to decode PNG: {path}", hint="file may be corrupt or zero-byte."
245            )
246        # cv2 returns BGR; BT.601 weights stay the same when applied to channels.
247        b, g, r = arr[..., 0], arr[..., 1], arr[..., 2]
248        return (0.114 * b + 0.587 * g + 0.299 * r).astype(np.uint8)
249    img = np.asarray(Image.open(path).convert("RGB"), dtype=np.float32)
250    luma = 0.299 * img[..., 0] + 0.587 * img[..., 1] + 0.114 * img[..., 2]
251    return luma.astype(np.uint8)
252
253
254def _load_pose(path: Path) -> Tuple[np.ndarray, np.ndarray]:
255    """T265 pose: ``position`` (3,), ``orientation`` (4,) in [x,y,z,w]."""
256    with np.load(path) as f:
257        pos = np.asarray(f["position"], dtype=np.float64).reshape(3)
258        quat = np.asarray(f["orientation"], dtype=np.float64).reshape(4)
259    return pos, quat
260
261
262# --------------------------------------------------------------------------- #
263# Annotation-effort extractor (no per-frame I/O — reads small txt files only)
264# --------------------------------------------------------------------------- #
265
266
267def _annotation_effort(phase_dir: Path, frame_stems: List[str]) -> Dict[str, float]:
268    """``iter-mean`` / ``iter-max``.
269
270    ``k_j = 1 + |{ i : frame j ∈ iter_i_faulty.txt }|``. A frame absent from
271    every iter_i_faulty.txt was accepted on the first pass (k=1); a frame
272    listed in all four was still faulty after iter 4 (k=5).
273    """
274    sam_dir = phase_dir / "sam2"
275    counts = np.ones(len(frame_stems), dtype=np.int32)
276    if not sam_dir.is_dir() or counts.size == 0:
277        return {"iter_mean": 1.0, "iter_max": 1.0}
278
279    stem_to_idx = {s: i for i, s in enumerate(frame_stems)}
280    for i in (1, 2, 3, 4):
281        f = sam_dir / f"iter{i}_faulty.txt"
282        if not f.is_file():
283            continue
284        for raw in f.read_text().splitlines():
285            entry = raw.strip()
286            if not entry:
287                continue
288            stem = Path(entry).stem  # tolerates "00007", "00007.png", "/abs/.../00007.png"
289            idx = stem_to_idx.get(stem)
290            if idx is not None:
291                counts[idx] += 1
292
293    return {"iter_mean": float(counts.mean()), "iter_max": float(counts.max())}
294
295
296# --------------------------------------------------------------------------- #
297# Geometry helpers
298# --------------------------------------------------------------------------- #
299
300
301def _bbox_xywh(binary_mask: np.ndarray) -> Optional[Tuple[int, int, int, int]]:
302    """Tight bbox of a boolean mask, or None if empty."""
303    rows = np.flatnonzero(binary_mask.any(axis=1))
304    cols = np.flatnonzero(binary_mask.any(axis=0))
305    if rows.size == 0 or cols.size == 0:
306        return None
307    y0, y1 = int(rows[0]), int(rows[-1])
308    x0, x1 = int(cols[0]), int(cols[-1])
309    return x0, y0, x1 - x0 + 1, y1 - y0 + 1
310
311
312def _occlusion_from_per_frame_bboxes(
313    bboxes_per_frame: List[Dict[int, Tuple[int, int, int, int]]],
314) -> Dict[str, float]:
315    """``occ-mean`` / ``occ-p90`` / ``occ-heavy`` from per-frame bbox dicts.
316
317    Per-object occlusion = mean over frames (where the object is visible) of
318    that frame's bbox-overlap-fraction. Then aggregate over objects.
319    """
320    per_object: Dict[int, List[float]] = {}
321    for bboxes in bboxes_per_frame:
322        if not bboxes:
323            continue
324        if len(bboxes) < 2:
325            for i in bboxes:
326                per_object.setdefault(i, []).append(0.0)
327            continue
328        for i, (xi, yi, wi, hi) in bboxes.items():
329            area_i = wi * hi
330            if area_i == 0:
331                continue
332            inter_sum = 0
333            for j, (xj, yj, wj, hj) in bboxes.items():
334                if i == j:
335                    continue
336                ix0, iy0 = max(xi, xj), max(yi, yj)
337                ix1, iy1 = min(xi + wi, xj + wj), min(yi + hi, yj + hj)
338                if ix1 > ix0 and iy1 > iy0:
339                    inter_sum += (ix1 - ix0) * (iy1 - iy0)
340            per_object.setdefault(i, []).append(inter_sum / area_i)
341
342    if not per_object:
343        return {"occ_mean": 0.0, "occ_p90": 0.0, "occ_heavy": 0.0}
344
345    per_object_mean = np.asarray(
346        [float(np.mean(v)) for v in per_object.values()],
347        dtype=np.float32,
348    )
349    return {
350        "occ_mean": float(per_object_mean.mean()),
351        "occ_p90": float(np.percentile(per_object_mean, 90)),
352        "occ_heavy": float((per_object_mean > _OCC_HEAVY_THRESHOLD).mean()),
353    }
354
355
356def _temporal_stability_from_areas(areas_T_by_id: np.ndarray) -> Dict[str, float]:
357    """``area-cv`` / ``area-drop`` / ``vis-instability`` from a (T, n_ids) area matrix.
358
359    Column ``j`` is the per-frame pixel count of instance ID ``j+1``.
360
361    ``area_drop`` follows old paper §3.4 eq 15: per-instance fraction of frames
362    where ``(A_{t-1} - A_t) / A_{t-1} > δ`` (with ``A_{t-1} > 0``), then
363    averaged over instances. Captures sudden mask collapses (typically caused
364    by an arm sweeping across the object during interaction).
365    """
366    T, n_ids = areas_T_by_id.shape
367    if n_ids == 0 or T == 0:
368        return {"area_cv": 0.0, "area_drop": 0.0, "vis_instability": 0.0}
369
370    present = areas_T_by_id > 0
371    cvs, drops, instabilities = [], [], []
372    for j in range(n_ids):
373        col = areas_T_by_id[:, j]
374        pres_col = present[:, j]
375
376        if not pres_col.any():
377            continue  # instance never observed (column would be all zeros)
378
379        positive = col[pres_col]
380        mean_a = float(positive.mean())
381        if mean_a > 0:
382            cvs.append(float(positive.std() / mean_a))
383
384        if T > 1:
385            flips = int(np.sum(pres_col[1:] != pres_col[:-1]))
386            instabilities.append(flips / T)
387
388            prev = col[:-1].astype(np.float64)
389            curr = col[1:].astype(np.float64)
390            valid = prev > 0
391            if valid.any():
392                ratios = np.zeros_like(prev)
393                ratios[valid] = (prev[valid] - curr[valid]) / prev[valid]
394                n_drop = int(np.sum(ratios > _AREA_DROP_THRESHOLD))
395                drops.append(n_drop / T)
396
397    return {
398        "area_cv": float(np.mean(cvs)) if cvs else 0.0,
399        "area_drop": float(np.mean(drops)) if drops else 0.0,
400        "vis_instability": float(np.mean(instabilities)) if instabilities else 0.0,
401    }
402
403
404def _quat_angle(q1: np.ndarray, q2: np.ndarray) -> float:
405    """Geodesic angle (rad) between two unit quaternions [x,y,z,w]."""
406    n1 = float(np.linalg.norm(q1))
407    n2 = float(np.linalg.norm(q2))
408    if n1 == 0 or n2 == 0:
409        return 0.0
410    dot = abs(float(np.dot(q1 / n1, q2 / n2)))
411    if dot > 1.0:
412        dot = 1.0
413    return float(2.0 * np.arccos(dot))
414
415
416def _mad_filter(arr: np.ndarray, k: float = 5.0) -> np.ndarray:
417    """Median-absolute-deviation outlier filter.
418
419    Replaces values more than ``k`` MADs from the median with the median.
420    This removes T265 pose jumps (relocalization spikes) that would
421    otherwise dominate jerk and trans_p90.
422    """
423    if arr.size < 3:
424        return arr
425    med = float(np.median(arr))
426    mad = float(np.median(np.abs(arr - med)))
427    if mad == 0:
428        return arr  # constant or near-constant signal
429    threshold = med + k * mad * 1.4826  # 1.4826 ≈ consistency constant for normal
430    out = arr.copy()
431    out[out > threshold] = med
432    return out
433
434
435def _camera_motion_from_arrays(
436    positions: np.ndarray,
437    quats: np.ndarray,
438) -> Dict[str, float]:
439    """Camera motion features with MAD outlier filtering.
440
441    - ``trans_mean`` / ``trans_p90``: per-frame translation magnitude (m)
442    - ``rot_mean`` / ``rot_p90``: per-frame rotation angle (rad)
443    - ``jerk``: mean |Δtranslation_{t+1} - Δtranslation_t| (m/frame²)
444
445    T265 VIO can produce occasional pose jumps from tracking-loss events;
446    MAD filtering removes these before computing statistics so they don't
447    dominate the p90 and jerk features.
448    """
449    if positions.shape[0] < 2:
450        return {"trans_mean": 0.0, "trans_p90": 0.0, "rot_mean": 0.0, "rot_p90": 0.0, "jerk": 0.0}
451
452    delta_t = np.linalg.norm(np.diff(positions, axis=0), axis=1)
453    delta_r = np.asarray(
454        [_quat_angle(quats[i], quats[i + 1]) for i in range(quats.shape[0] - 1)],
455        dtype=np.float64,
456    )
457    # MAD-filter to remove T265 relocalization spikes.
458    delta_t = _mad_filter(delta_t)
459    delta_r = _mad_filter(delta_r)
460    jerk = np.abs(np.diff(delta_t)) if delta_t.size >= 2 else np.zeros(0)
461
462    return {
463        "trans_mean": float(delta_t.mean()),
464        "trans_p90": float(np.percentile(delta_t, 90)),
465        "rot_mean": float(delta_r.mean()),
466        "rot_p90": float(np.percentile(delta_r, 90)),
467        "jerk": float(jerk.mean()) if jerk.size else 0.0,
468    }
469
470
471# --------------------------------------------------------------------------- #
472# Fisheye / stereo features (old paper §3.7 eqs 25-29)
473# --------------------------------------------------------------------------- #
474
475_FISHEYE_FEATURE_NAMES: Tuple[str, ...] = (
476    "fisheye_dark",
477    "fisheye_bright",
478    "fisheye_sharpness",
479    "fisheye_corr",
480    "fisheye_texture",
481)
482# Regex for ``*_L.png`` / ``*_R.png`` style fisheye pairs.
483_FISHEYE_LR_RE = re.compile(r"^(?P<stem>.+?)[_-](?P<side>[LRlr])$")
484
485
486def _discover_fisheye(
487    fisheye_dir: Path,
488) -> Tuple[List[str], Dict[str, Path], Dict[str, Path]]:
489    """Find fisheye frames; return (sorted_stems, left_path_map, right_path_map).
490
491    Auto-detects three layouts in priority order:
492
493    1. Subdirs ``left/`` + ``right/`` with matching basenames.
494    2. Files like ``00000_L.png`` + ``00000_R.png`` (or ``-L``/``-R``).
495    3. Single PNGs in the dir → treated as left only; ``fisheye_corr`` will
496       be unset for those frames.
497
498    Returns ``([], {}, {})`` if the directory is missing or empty.
499    """
500    if not fisheye_dir.is_dir():
501        return [], {}, {}
502
503    # Layout A: left/ + right/ subdirectories.
504    ld, rd = fisheye_dir / "left", fisheye_dir / "right"
505    if ld.is_dir() and rd.is_dir():
506        left = {p.stem: p for p in ld.iterdir() if p.suffix.lower() in (".png", ".jpg")}
507        right = {p.stem: p for p in rd.iterdir() if p.suffix.lower() in (".png", ".jpg")}
508        common = sorted(set(left) & set(right))
509        return common, left, right
510
511    # Layout B: <stem>_L.png + <stem>_R.png.
512    files = [
513        p for p in fisheye_dir.iterdir() if p.is_file() and p.suffix.lower() in (".png", ".jpg")
514    ]
515    left, right = {}, {}
516    for f in files:
517        m = _FISHEYE_LR_RE.match(f.stem)
518        if m:
519            stem = m.group("stem")
520            if m.group("side").lower() == "l":
521                left[stem] = f
522            else:
523                right[stem] = f
524    if left and right:
525        common = sorted(set(left) & set(right))
526        return common, left, right
527
528    # Layout C: single fisheye images (no stereo pair). corr → undefined.
529    singles = {p.stem: p for p in files}
530    if singles:
531        return sorted(singles), singles, {}
532
533    return [], {}, {}
534
535
536def _laplacian_var(img: np.ndarray) -> float:
537    """Variance of Laplacian — image sharpness / blur measure."""
538    if _HAS_CV2:
539        lap = cv2.Laplacian(img, cv2.CV_32F)
540        return float(lap.var())
541    f = img.astype(np.float32)
542    lap = np.roll(f, 1, 0) + np.roll(f, -1, 0) + np.roll(f, 1, 1) + np.roll(f, -1, 1) - 4.0 * f
543    return float(lap.var())
544
545
546def _gradient_magnitude_mean(img: np.ndarray) -> float:
547    """Mean of pixelwise |∇F| — texture density."""
548    f = img.astype(np.float32)
549    gy, gx = np.gradient(f)
550    return float(np.sqrt(gx * gx + gy * gy).mean())
551
552
553def _pearson_corr_2d(a: np.ndarray, b: np.ndarray) -> Optional[float]:
554    """Pearson correlation between two same-size single-channel images.
555
556    Resizes ``b`` to ``a``'s shape if they differ — matches the literal paper
557    formula ``corr(F^L, F^R)`` without requiring stereo rectification.
558    """
559    if a.shape != b.shape:
560        if _HAS_CV2:
561            b = cv2.resize(b, (a.shape[1], a.shape[0]))
562        else:
563            b = np.asarray(Image.fromarray(b).resize((a.shape[1], a.shape[0])))
564    af = a.astype(np.float64).ravel()
565    bf = b.astype(np.float64).ravel()
566    af -= af.mean()
567    bf -= bf.mean()
568    denom = math.sqrt((af * af).sum() * (bf * bf).sum())
569    if denom == 0:
570        return None  # constant image — correlation undefined
571    return float((af * bf).sum() / denom)
572
573
574def _fisheye_features(phase_dir: Path) -> Dict[str, float]:
575    """Compute the 5 fisheye features for one (scene, phase).
576
577    Returns all-NaN dict if no fisheye dir / no decodable frames found,
578    so that percentile normalization can exclude missing scenes instead of
579    biasing them toward low difficulty.
580    Logs a warning when stereo pairs are missing (``fisheye_corr`` then
581    aggregates only over frames that DO have pairs; if none, → NaN).
582    """
583    empty = {name: float("nan") for name in _FISHEYE_FEATURE_NAMES}
584    stems, left_paths, right_paths = _discover_fisheye(phase_dir / "fisheye")
585    if not stems:
586        return empty
587
588    has_right = bool(right_paths)
589    if not has_right:
590        log.warning(
591            "%s/fisheye: no stereo pairs found; fisheye_corr will be 0",
592            phase_dir,
593        )
594
595    dark_acc: List[float] = []
596    bright_acc: List[float] = []
597    sharp_acc: List[float] = []
598    tex_acc: List[float] = []
599    corr_acc: List[float] = []
600
601    for s in stems:
602        try:
603            L = _decode_png_gray(left_paths[s])
604        except Exception as e:
605            log.warning("fisheye decode failed for %s: %s", left_paths[s], e)
606            continue
607
608        dark_acc.append(float((L < _DARK_LUMA_MAX).mean()))
609        bright_acc.append(float((L > _SPECULAR_LUMA_MIN).mean()))
610        sharp_acc.append(_laplacian_var(L))
611        tex_acc.append(_gradient_magnitude_mean(L))
612
613        if has_right and s in right_paths:
614            try:
615                R = _decode_png_gray(right_paths[s])
616                c = _pearson_corr_2d(L, R)
617                if c is not None:
618                    corr_acc.append(c)
619            except Exception as e:
620                log.warning("fisheye-R decode failed for %s: %s", right_paths[s], e)
621
622    return {
623        "fisheye_dark": float(np.mean(dark_acc)) if dark_acc else 0.0,
624        "fisheye_bright": float(np.mean(bright_acc)) if bright_acc else 0.0,
625        "fisheye_sharpness": float(np.mean(sharp_acc)) if sharp_acc else 0.0,
626        "fisheye_corr": float(np.mean(corr_acc)) if corr_acc else 0.0,
627        "fisheye_texture": float(np.mean(tex_acc)) if tex_acc else 0.0,
628    }
629
630
631# --------------------------------------------------------------------------- #
632# Top-level driver — streaming, single I/O pass
633# --------------------------------------------------------------------------- #
634
635
636def _scene_phase_from_dir(phase_dir: Path) -> Tuple[str, int]:
637    scene_id = phase_dir.parent.name
638    try:
639        phase_idx = int(phase_dir.name)
640    except ValueError as e:
641        raise DatasetError(
642            f"phase dir name must be an integer, got {phase_dir.name!r}",
643            hint="phase_dir is expected to be <data_root>/<scene>/<phase_idx>; "
644            "rename or pass a deeper path.",
645        ) from e
646    return scene_id, phase_idx
647
648
649# Features that should be NaN (not 0) when their modality is absent.
650_NAN_WHEN_ABSENT: frozenset = frozenset(
651    {
652        "fisheye_dark",
653        "fisheye_bright",
654        "fisheye_sharpness",
655        "fisheye_corr",
656        "fisheye_texture",
657    }
658)
659
660
661def _empty_features() -> Dict[str, float]:
662    return {k: (float("nan") if k in _NAN_WHEN_ABSENT else 0.0) for k in FEATURE_NAMES}
663
664
665def extract_phase_features(phase_dir: Path) -> PhaseFeatures:
666    """Compute all 18 ESD features for one (scene, phase) directory.
667
668    ``phase_dir`` is expected to be ``<data_root>/<scene_dir>/<phase_idx>``;
669    e.g. ``test_dataset_aggregated/scene1.GDC.cse/0``. The scene_id and
670    phase fields of the result are inferred from this path.
671
672    The per-frame loop is streaming: one frame's modalities are decoded,
673    fed into per-feature accumulators, and discarded. Peak memory is
674    independent of phase length.
675    """
676    phase_dir = Path(phase_dir)
677    scene_id, phase_idx = _scene_phase_from_dir(phase_dir)
678
679    rgb_dir = phase_dir / "rgb"
680    n_total = len(_stem_set(rgb_dir, ".png")) if rgb_dir.is_dir() else 0
681    stems = _aligned_frames(phase_dir)
682
683    if not stems:
684        log.warning("no aligned frames found for %s/%d", scene_id, phase_idx)
685        return PhaseFeatures(
686            scene_id=scene_id,
687            phase=phase_idx,
688            n_frames_total=n_total,
689            n_frames_used=0,
690            features=_empty_features(),
691        )
692
693    T = len(stems)
694
695    # Annotation-effort features need no per-frame I/O.
696    feats: Dict[str, float] = dict(_annotation_effort(phase_dir, stems))
697
698    # Streaming accumulators.
699    obj_counts = np.zeros(T, dtype=np.int32)
700    depth_invalid_frac = np.zeros(T, dtype=np.float64)
701    depth_inv_mask_frac = np.full(T, np.nan, dtype=np.float64)  # NaN → no mask pixels
702    depth_std_per_frame = np.full(T, np.nan, dtype=np.float64)  # NaN → all-invalid
703    depth_std_mask_per_frame = np.full(T, np.nan, dtype=np.float64)  # NaN → no valid depth in mask
704    spec_frac = np.zeros(T, dtype=np.float64)
705    dark_frac = np.zeros(T, dtype=np.float64)
706    rgb_blur_per_frame = np.zeros(T, dtype=np.float64)  # Laplacian variance
707    rgb_tex_per_frame = np.zeros(T, dtype=np.float64)  # gradient magnitude mean
708    bboxes_per_frame: List[Dict[int, Tuple[int, int, int, int]]] = []
709    # Per-frame dense bincount of mask values (length varies by frame max ID);
710    # zipped at the end into a (T, max_id) matrix for area_cv / vis_instability.
711    bincount_per_frame: List[np.ndarray] = []
712    positions = np.zeros((T, 3), dtype=np.float64)
713    quats = np.zeros((T, 4), dtype=np.float64)
714
715    for t, s in enumerate(stems):
716        depth_m = _load_depth_m(phase_dir / "depth" / f"{s}.png")
717        mask = _load_mask(phase_dir / "sam2/masks" / f"{s}.png")
718        rgb_luma = _load_rgb_luma(phase_dir / "rgb" / f"{s}.png")
719        pos, q = _load_pose(phase_dir / "cam_pose" / f"{s}.npz")
720
721        # ── Depth-quality features ────────────────────────────────────────
722        invalid = depth_m <= 0
723        depth_invalid_frac[t] = float(invalid.mean())
724
725        in_mask = mask > 0
726        if in_mask.any():
727            depth_inv_mask_frac[t] = float(invalid[in_mask].mean())
728
729        valid = depth_m[~invalid]
730        if valid.size > 0:
731            depth_std_per_frame[t] = float(valid.std())
732
733        # In-mask depth std: only valid-depth pixels that also fall inside any
734        # instance mask. Captures object-relevant depth noise (paper's
735        # depth_std is whole-image and dominated by background).
736        valid_in_mask = depth_m[(~invalid) & in_mask]
737        if valid_in_mask.size > 0:
738            depth_std_mask_per_frame[t] = float(valid_in_mask.std())
739
740        # ── Photometric-conflict features ─────────────────────────────────
741        spec_frac[t] = float(((rgb_luma > _SPECULAR_LUMA_MIN) & invalid).mean())
742        dark_frac[t] = float(((rgb_luma < _DARK_LUMA_MAX) & invalid).mean())
743
744        # ── Image quality (D435 RGB) — grounded in CEMS (Zhao et al. 2024) ──
745        rgb_blur_per_frame[t] = _laplacian_var(rgb_luma)
746        rgb_tex_per_frame[t] = _gradient_magnitude_mean(rgb_luma)
747
748        # ── Mask-derived features (one bincount per frame, reused 3×) ─────
749        flat = mask.ravel()
750        # Drop negative IDs (defensive — masks should be unsigned).
751        if flat.min() < 0:
752            flat = flat[flat >= 0]
753        counts = np.bincount(flat)  # counts[i] = pixel count for ID i
754        ids = np.flatnonzero(counts[1:]) + 1 if counts.size > 1 else np.empty(0, dtype=np.int64)
755        obj_counts[t] = ids.size
756        bincount_per_frame.append(counts)
757
758        bboxes: Dict[int, Tuple[int, int, int, int]] = {}
759        for oid in ids:
760            b = _bbox_xywh(mask == oid)
761            if b is not None:
762                bboxes[int(oid)] = b
763        bboxes_per_frame.append(bboxes)
764
765        # ── Pose ──────────────────────────────────────────────────────────
766        positions[t] = pos
767        quats[t] = q
768
769    # ── n_total_objects: prefer mask_to_object.json, else union of observed ──
770    obj_json = phase_dir / "sam2" / "mask_to_object.json"
771    if obj_json.is_file():
772        try:
773            n_total_objects = len(json.loads(obj_json.read_text()))
774        except (ValueError, OSError) as e:
775            log.warning("%s: cannot parse, falling back to observed IDs (%s)", obj_json, e)
776            n_total_objects = max(
777                (int(c.size - 1) for c in bincount_per_frame if c.size > 1),
778                default=0,
779            )
780    else:
781        all_ids = {
782            int(i)
783            for c in bincount_per_frame
784            for i in (np.flatnonzero(c[1:]) + 1 if c.size > 1 else ())
785        }
786        n_total_objects = len(all_ids)
787
788    # ── Scene-complexity features ────────────────────────────────────────
789    feats["obj_mean"] = float(obj_counts.mean())
790    feats["obj_std"] = float(obj_counts.std())
791    feats["obj_consist"] = (
792        float((obj_counts == n_total_objects).mean()) if n_total_objects > 0 else 0.0
793    )
794
795    # ── Occlusion ────────────────────────────────────────────────────────
796    feats.update(_occlusion_from_per_frame_bboxes(bboxes_per_frame))
797
798    # ── Depth aggregates (NaN means "skip frame for this aggregate") ─────
799    feats["depth_invalid"] = float(depth_invalid_frac.mean())
800    feats["depth_invalid_mask"] = (
801        float(np.nanmean(depth_inv_mask_frac)) if not np.isnan(depth_inv_mask_frac).all() else 0.0
802    )
803    feats["depth_std"] = (
804        float(np.nanmean(depth_std_per_frame)) if not np.isnan(depth_std_per_frame).all() else 0.0
805    )
806    feats["depth_std_mask"] = (
807        float(np.nanmean(depth_std_mask_per_frame))
808        if not np.isnan(depth_std_mask_per_frame).all()
809        else 0.0
810    )
811
812    # ── Photometric ──────────────────────────────────────────────────────
813    feats["specular"] = float(spec_frac.mean())
814    feats["dark"] = float(dark_frac.mean())
815
816    # ── Image quality (D435 RGB) ─────────────────────────────────────
817    feats["rgb_blur"] = float(rgb_blur_per_frame.mean())
818    feats["rgb_texture"] = float(rgb_tex_per_frame.mean())
819
820    # ── Compute max_id once for reuse by object-size and temporal-stability ─
821    max_id = max((int(c.size - 1) for c in bincount_per_frame if c.size > 1), default=0)
822
823    # ── Object size distribution ────────────────────────────────────
824    # Per-instance mean mask area fraction (over frames where visible),
825    # then mean/std across instances. Small objects are universally harder.
826    if max_id >= 1 and T > 0:
827        _frame_pixels = bincount_per_frame[0].sum() if bincount_per_frame else 1
828        _areas_for_size = np.zeros((T, max_id), dtype=np.int64)
829        for _t, _c in enumerate(bincount_per_frame):
830            _n = min(int(_c.size) - 1, max_id)
831            if _n > 0:
832                _areas_for_size[_t, :_n] = _c[1 : 1 + _n]
833        _present = _areas_for_size > 0
834        _per_inst_mean = []
835        for _j in range(max_id):
836            _vis = _areas_for_size[:, _j][_present[:, _j]]
837            if _vis.size > 0:
838                _per_inst_mean.append(float(_vis.mean()) / max(_frame_pixels, 1))
839        if _per_inst_mean:
840            feats["mask_area_mean"] = float(np.mean(_per_inst_mean))
841            feats["mask_area_std"] = float(np.std(_per_inst_mean))
842        else:
843            feats["mask_area_mean"] = 0.0
844            feats["mask_area_std"] = 0.0
845    else:
846        feats["mask_area_mean"] = 0.0
847        feats["mask_area_std"] = 0.0
848
849    # ── Temporal stability: build dense (T, max_id) area matrix ──────────
850    if max_id >= 1:
851        areas = np.zeros((T, max_id), dtype=np.int64)
852        for t, c in enumerate(bincount_per_frame):
853            n = min(int(c.size) - 1, max_id)
854            if n > 0:
855                areas[t, :n] = c[1 : 1 + n]
856        feats.update(_temporal_stability_from_areas(areas))
857    else:
858        feats["area_cv"] = 0.0
859        feats["vis_instability"] = 0.0
860
861    # ── Camera motion ────────────────────────────────────────────────────
862    feats.update(_camera_motion_from_arrays(positions, quats))
863
864    # ── Fisheye / stereo (independent I/O — uses fisheye/, not the
865    # rgb/depth/mask/pose modalities streamed above) ─────────────────────
866    feats.update(_fisheye_features(phase_dir))
867
868    # Stable column order matching FEATURE_NAMES.
869    # Missing fisheye features → NaN (not 0) so percentile normalization
870    # can exclude them instead of biasing toward low difficulty.
871    feats = {
872        k: float(feats.get(k, float("nan") if k in _NAN_WHEN_ABSENT else 0.0))
873        for k in FEATURE_NAMES
874    }
875
876    return PhaseFeatures(
877        scene_id=scene_id,
878        phase=phase_idx,
879        n_frames_total=n_total,
880        n_frames_used=T,
881        features=feats,
882    )
883
884
885def iter_phase_dirs(data_root: Path) -> Iterator[Path]:
886    """Yield every ``<scene>/<phase>`` directory under ``data_root``.
887
888    A phase directory is any immediate child of a scene directory whose
889    name is purely numeric and that contains an ``rgb/`` subdirectory.
890    """
891    data_root = Path(data_root)
892    if not data_root.is_dir():
893        raise DatasetError(
894            f"data_root is not a directory: {data_root}", hint="check the --data-root argument."
895        )
896    for scene_dir in sorted(p for p in data_root.iterdir() if p.is_dir()):
897        for child in sorted(scene_dir.iterdir()):
898            if child.is_dir() and child.name.isdigit() and (child / "rgb").is_dir():
899                yield child