scripts/data/audit_trex_tracks.py

#!/usr/bin/env python3
"""Audit all T-Rex track caches and investigate near-static wrist tracks."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import cv2
import numpy as np
import pyarrow.parquet as pq

VIEW_SLICES = {
    "head_left": (0, 100),
    "left_wrist": (100, 175),
    "right_wrist": (175, 250),
}
WRIST_GROUPS = {
    "left_wrist": {"background": (100, 125), "hand": (125, 175)},
    "right_wrist": {"background": (175, 200), "hand": (200, 250)},
}
IMAGE_SCALE = np.array([320.0, 180.0], dtype=np.float32)
WINDOW_FRAMES = 768
WINDOW_OVERLAP = 64
WINDOW_STEP = WINDOW_FRAMES - WINDOW_OVERLAP


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--dataset-root",
        type=Path,
        default=Path("/scratch1/home/zhicao/dreamzero/data/trex_full_force"),
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path(
            "/scratch1/home/zhicao/dreamzero/data/"
            "trex_full_force/audit/track_quality"
        ),
    )
    parser.add_argument("--static-span-px", type=float, default=2.0)
    parser.add_argument("--video-samples", type=int, default=5)
    parser.add_argument("--video-workers", type=int, default=8)
    parser.add_argument("--skip-video-check", action="store_true")
    return parser.parse_args()


def _quantiles(values: list[float]) -> dict[str, float]:
    array = np.asarray(values, dtype=np.float64)
    return {
        name: float(np.percentile(array, percentile))
        for name, percentile in (
            ("min", 0),
            ("p01", 1),
            ("p05", 5),
            ("median", 50),
            ("p95", 95),
            ("p99", 99),
            ("max", 100),
        )
    }


def _point_span_px(tracks: np.ndarray) -> np.ndarray:
    pixel_tracks = np.asarray(tracks, dtype=np.float32) * IMAGE_SCALE
    return np.sqrt(
        np.ptp(pixel_tracks[..., 0], axis=0) ** 2
        + np.ptp(pixel_tracks[..., 1], axis=0) ** 2
    )


def _seam_metrics(
    tracks: np.ndarray,
    visibility: np.ndarray,
) -> tuple[float, float, float]:
    frames = int(tracks.shape[0])
    boundaries = list(range(WINDOW_FRAMES, frames, WINDOW_STEP))
    if not boundaries:
        return math.nan, math.nan, math.nan

    pixel_tracks = np.asarray(tracks, dtype=np.float32) * IMAGE_SCALE
    delta = np.linalg.norm(np.diff(pixel_tracks, axis=0), axis=-1)
    visible_pair = (visibility[1:] > 0.5) & (visibility[:-1] > 0.5)
    seam_indices = np.asarray([boundary - 1 for boundary in boundaries], dtype=np.int64)
    seam_values = delta[seam_indices][visible_pair[seam_indices]]

    regular_indices = np.unique(
        np.rint(np.linspace(0, max(0, frames - 2), min(128, frames - 1))).astype(
            np.int64
        )
    )
    regular_indices = regular_indices[
        ~np.isin(regular_indices, seam_indices)
    ]
    regular_values = delta[regular_indices][visible_pair[regular_indices]]
    seam_p95 = (
        float(np.percentile(seam_values, 95)) if seam_values.size else math.nan
    )
    regular_p95 = (
        float(np.percentile(regular_values, 95))
        if regular_values.size
        else math.nan
    )
    ratio = (
        seam_p95 / max(regular_p95, 0.1)
        if np.isfinite(seam_p95) and np.isfinite(regular_p95)
        else math.nan
    )
    return seam_p95, regular_p95, ratio


def _audit_track(
    path: Path,
    episode_index: int,
    expected_frames: int,
) -> tuple[dict[str, object], list[str]]:
    errors: list[str] = []
    with np.load(path, allow_pickle=False) as payload:
        tracks = np.asarray(payload["tracks"], dtype=np.float32)
        visibility = np.asarray(payload["vis"], dtype=np.float32)
    if tracks.shape != (expected_frames, 250, 2):
        errors.append(f"tracks shape {tracks.shape} != {(expected_frames, 250, 2)}")
    if visibility.shape != (expected_frames, 250):
        errors.append(
            f"visibility shape {visibility.shape} != {(expected_frames, 250)}"
        )
    if errors:
        return {"episode_index": episode_index, "frames": expected_frames}, errors

    finite = np.isfinite(tracks).all(axis=-1)
    in_frame = ((tracks >= 0.0) & (tracks <= 1.0)).all(axis=-1)
    binary_visibility = (visibility == 0.0) | (visibility == 1.0)
    if not finite.all():
        errors.append("non-finite track coordinates")
    if not in_frame.all():
        errors.append("out-of-range normalized track coordinates")
    if not binary_visibility.all():
        errors.append("non-binary visibility")

    row: dict[str, object] = {
        "episode_index": episode_index,
        "frames": expected_frames,
        "finite_fraction": float(finite.mean()),
        "in_frame_fraction": float(in_frame.mean()),
        "binary_visibility_fraction": float(binary_visibility.mean()),
    }
    for view, (start, end) in VIEW_SLICES.items():
        view_tracks = tracks[:, start:end]
        view_visibility = visibility[:, start:end]
        span = _point_span_px(view_tracks)
        row[f"{view}_visibility"] = float(view_visibility.mean())
        row[f"{view}_span_median_px"] = float(np.median(span))
        row[f"{view}_span_p95_px"] = float(np.percentile(span, 95))
        row[f"{view}_static_point_fraction"] = float(np.mean(span < 1.0))
        if episode_index >= 1737:
            seam_p95, regular_p95, seam_ratio = _seam_metrics(
                view_tracks,
                view_visibility,
            )
        else:
            seam_p95, regular_p95, seam_ratio = math.nan, math.nan, math.nan
        row[f"{view}_seam_jump_p95_px"] = seam_p95
        row[f"{view}_regular_jump_p95_px"] = regular_p95
        row[f"{view}_seam_jump_ratio"] = seam_ratio

    for view, groups in WRIST_GROUPS.items():
        for group, (start, end) in groups.items():
            span = _point_span_px(tracks[:, start:end])
            row[f"{view}_{group}_span_median_px"] = float(np.median(span))
            row[f"{view}_{group}_span_p95_px"] = float(
                np.percentile(span, 95)
            )
    return row, errors


def _video_path(root: Path, episode_index: int, view: str) -> Path:
    return (
        root
        / "videos"
        / f"chunk-{episode_index // 1000:03d}"
        / f"observation.images.{view}"
        / f"episode_{episode_index:06d}.mp4"
    )


def _sample_video_motion(
    root: Path,
    episode_index: int,
    view: str,
    sample_count: int,
) -> dict[str, object]:
    path = _video_path(root, episode_index, view)
    cap = cv2.VideoCapture(str(path))
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    frames: list[np.ndarray] = []
    if frame_count > 0:
        sample_indices = np.rint(
            np.linspace(0, frame_count - 1, sample_count)
        ).astype(int)
        for frame_index in sample_indices:
            cap.set(cv2.CAP_PROP_POS_FRAMES, int(frame_index))
            ok, frame = cap.read()
            if not ok:
                continue
            gray = cv2.resize(
                cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY),
                (80, 45),
                interpolation=cv2.INTER_AREA,
            ).astype(np.float32)
            frames.append(gray)
    cap.release()
    if len(frames) < 2:
        return {
            "episode_index": episode_index,
            "view": view,
            "video_error": f"decoded only {len(frames)} sampled frame(s)",
        }
    adjacent = [
        float(np.mean(np.abs(second - first)))
        for first, second in zip(frames, frames[1:])
    ]
    return {
        "episode_index": episode_index,
        "view": view,
        "video_frames": frame_count,
        "video_adjacent_mad_mean": float(np.mean(adjacent)),
        "video_adjacent_mad_max": float(np.max(adjacent)),
        "video_first_last_mad": float(np.mean(np.abs(frames[-1] - frames[0]))),
        "video_first_frame_hash": hashlib.sha256(
            frames[0].astype(np.uint8).tobytes()
        ).hexdigest()[:16],
        "video_sample_hash": hashlib.sha256(
            np.stack(frames).astype(np.uint8).tobytes()
        ).hexdigest()[:16],
    }


def _arm_motion(root: Path, episode_index: int) -> dict[str, float]:
    path = (
        root
        / "data"
        / f"chunk-{episode_index // 1000:03d}"
        / f"episode_{episode_index:06d}.parquet"
    )
    column = pq.read_table(path, columns=["observation.state"]).column(0)
    array = column.combine_chunks()
    values = np.asarray(array.values.to_numpy(zero_copy_only=False)).reshape(
        len(array),
        -1,
    )
    result: dict[str, float] = {}
    for name, selection in (
        ("left_arm", slice(0, 7)),
        ("right_arm", slice(29, 36)),
    ):
        arm = values[:, selection].astype(np.float32, copy=False)
        result[f"{name}_joint_path_l2"] = float(
            np.linalg.norm(np.diff(arm, axis=0), axis=1).sum()
        )
        result[f"{name}_end_delta_l2"] = float(
            np.linalg.norm(arm[-1] - arm[0])
        )
        result[f"{name}_max_joint_range"] = float(
            np.ptp(arm, axis=0).max()
        )
    return result


def main() -> int:
    args = _parse_args()
    root = args.dataset_root.expanduser().resolve()
    output_dir = args.output_dir.expanduser().resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    episodes = [
        json.loads(line)
        for line in (root / "meta" / "episodes.jsonl").read_text().splitlines()
        if line.strip()
    ]
    rows: list[dict[str, object]] = []
    errors: list[dict[str, object]] = []
    for position, episode in enumerate(episodes, start=1):
        episode_index = int(episode["episode_index"])
        path = (
            root
            / "tracks_trex_track_force_v2"
            / f"episode_{episode_index:06d}.npz"
        )
        try:
            row, track_errors = _audit_track(
                path,
                episode_index,
                int(episode["length"]),
            )
        except Exception as exc:  # noqa: BLE001
            row = {
                "episode_index": episode_index,
                "frames": int(episode["length"]),
            }
            track_errors = [f"{type(exc).__name__}: {exc}"]
        row["task"] = " | ".join(episode.get("tasks", []))
        rows.append(row)
        if track_errors:
            errors.append(
                {
                    "episode_index": episode_index,
                    "errors": track_errors,
                }
            )
        if position % 250 == 0 or position == len(episodes):
            print(f"Audited tracks: {position}/{len(episodes)}", flush=True)

    static_threshold = float(args.static_span_px)
    candidates: list[tuple[int, str]] = []
    for row in rows:
        for view in ("left_wrist", "right_wrist"):
            value = row.get(f"{view}_background_span_median_px")
            if isinstance(value, float) and value < static_threshold:
                candidates.append((int(row["episode_index"]), view))
    candidate_keys = set(candidates)
    candidate_episode_indices = sorted({episode for episode, _ in candidates})

    video_results: list[dict[str, object]] = []
    video_targets = [
        (int(row["episode_index"]), view)
        for row in rows
        for view in ("left_wrist", "right_wrist")
    ]
    if not args.skip_video_check:
        with ThreadPoolExecutor(max_workers=max(1, int(args.video_workers))) as pool:
            futures = {
                pool.submit(
                    _sample_video_motion,
                    root,
                    episode_index,
                    view,
                    max(2, int(args.video_samples)),
                ): (episode_index, view)
                for episode_index, view in video_targets
            }
            for position, future in enumerate(as_completed(futures), start=1):
                video_results.append(future.result())
                if position % 500 == 0 or position == len(futures):
                    print(
                        f"Checked wrist videos: "
                        f"{position}/{len(futures)}",
                        flush=True,
                    )

    video_by_key = {
        (int(result["episode_index"]), str(result["view"])): result
        for result in video_results
    }
    rows_by_episode = {
        int(row["episode_index"]): row
        for row in rows
    }
    video_frame_count_mismatches = [
        {
            "episode_index": int(result["episode_index"]),
            "view": str(result["view"]),
            "expected_frames": int(
                rows_by_episode[int(result["episode_index"])]["frames"]
            ),
            "actual_frames": int(result["video_frames"]),
        }
        for result in video_results
        if "video_frames" in result
        and int(result["video_frames"])
        != int(rows_by_episode[int(result["episode_index"])]["frames"])
    ]
    video_near_static_keys = {
        key
        for key, video in video_by_key.items()
        if not video.get("video_error")
        and float(video.get("video_adjacent_mad_mean", math.inf)) < 1.0
        and float(video.get("video_first_last_mad", math.inf)) < 2.0
    }
    video_error_keys = {
        key
        for key, video in video_by_key.items()
        if video.get("video_error")
    }
    relevant_keys = candidate_keys | video_near_static_keys | video_error_keys
    relevant_episode_indices = sorted(
        {episode for episode, _ in relevant_keys}
    )
    arm_by_episode: dict[int, dict[str, float]] = {}
    for position, episode_index in enumerate(
        relevant_episode_indices,
        start=1,
    ):
        arm_by_episode[episode_index] = _arm_motion(root, episode_index)
        if position % 250 == 0 or position == len(relevant_episode_indices):
            print(
                f"Checked relevant arm motion: "
                f"{position}/{len(relevant_episode_indices)}",
                flush=True,
            )

    classifications: dict[str, list[dict[str, object]]] = {
        "source_video_near_static_with_moving_arm": [],
        "track_near_static_with_moving_video_and_arm": [],
        "stationary_arm_or_low_motion": [],
        "video_check_error": [],
    }
    for episode_index, view in sorted(relevant_keys):
        row = rows[episode_index]
        video = video_by_key.get((episode_index, view), {})
        arm_name = "left_arm" if view == "left_wrist" else "right_arm"
        arm = arm_by_episode[episode_index]
        max_joint_range = float(arm[f"{arm_name}_max_joint_range"])
        video_error = video.get("video_error")
        key = (episode_index, view)
        video_near_static = key in video_near_static_keys
        track_near_static = key in candidate_keys
        arm_moving = max_joint_range >= 0.1
        record = {
            "episode_index": episode_index,
            "view": view,
            "frames": int(row["frames"]),
            "task": row["task"],
            "background_span_median_px": row[
                f"{view}_background_span_median_px"
            ],
            "hand_span_median_px": row[f"{view}_hand_span_median_px"],
            "video_adjacent_mad_mean": video.get(
                "video_adjacent_mad_mean"
            ),
            "video_first_last_mad": video.get("video_first_last_mad"),
            "video_first_frame_hash": video.get("video_first_frame_hash"),
            "video_sample_hash": video.get("video_sample_hash"),
            "arm_max_joint_range": max_joint_range,
            "track_near_static": track_near_static,
            "video_near_static": video_near_static,
        }
        if video_error:
            record["video_error"] = video_error
            classifications["video_check_error"].append(record)
        elif video_near_static and arm_moving:
            classifications[
                "source_video_near_static_with_moving_arm"
            ].append(record)
        elif track_near_static and not video_near_static and arm_moving:
            classifications[
                "track_near_static_with_moving_video_and_arm"
            ].append(record)
        else:
            classifications["stationary_arm_or_low_motion"].append(record)

    seam_anomalies: list[dict[str, object]] = []
    for row in rows:
        for view in VIEW_SLICES:
            ratio = row.get(f"{view}_seam_jump_ratio")
            seam_p95 = row.get(f"{view}_seam_jump_p95_px")
            if (
                isinstance(ratio, float)
                and isinstance(seam_p95, float)
                and np.isfinite(ratio)
                and np.isfinite(seam_p95)
                and ratio > 5.0
                and seam_p95 > 5.0
            ):
                seam_anomalies.append(
                    {
                        "episode_index": int(row["episode_index"]),
                        "view": view,
                        "seam_jump_p95_px": seam_p95,
                        "regular_jump_p95_px": row[
                            f"{view}_regular_jump_p95_px"
                        ],
                        "ratio": ratio,
                    }
                )

    metric_distributions: dict[str, dict[str, float]] = {}
    for view in VIEW_SLICES:
        for suffix in ("visibility", "span_median_px", "seam_jump_ratio"):
            key = f"{view}_{suffix}"
            values = [
                float(row[key])
                for row in rows
                if isinstance(row.get(key), float)
                and np.isfinite(float(row[key]))
            ]
            if values:
                metric_distributions[key] = _quantiles(values)
    for view in WRIST_GROUPS:
        for group in ("background", "hand"):
            key = f"{view}_{group}_span_median_px"
            metric_distributions[key] = _quantiles(
                [float(row[key]) for row in rows]
            )

    csv_path = output_dir / "episode_metrics.csv"
    fieldnames = sorted({key for row in rows for key in row})
    with csv_path.open("w", newline="") as file:
        writer = csv.DictWriter(file, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    video_csv_path = output_dir / "wrist_video_metrics.csv"
    if video_results:
        video_fieldnames = sorted(
            {key for result in video_results for key in result}
        )
        with video_csv_path.open("w", newline="") as file:
            writer = csv.DictWriter(file, fieldnames=video_fieldnames)
            writer.writeheader()
            writer.writerows(
                sorted(
                    video_results,
                    key=lambda result: (
                        int(result["episode_index"]),
                        str(result["view"]),
                    ),
                )
            )

    frozen_records = classifications[
        "source_video_near_static_with_moving_arm"
    ]
    frozen_episode_indices = sorted(
        {int(record["episode_index"]) for record in frozen_records}
    )
    frozen_views_by_episode: dict[int, set[str]] = {}
    for record in frozen_records:
        frozen_views_by_episode.setdefault(
            int(record["episode_index"]),
            set(),
        ).add(str(record["view"]))
    frozen_episode_breakdown = Counter(
        "both" if len(views) == 2 else next(iter(views))
        for views in frozen_views_by_episode.values()
    )
    frozen_frame_count = sum(
        int(rows_by_episode[episode_index]["frames"])
        for episode_index in frozen_episode_indices
    )
    first_frame_hash_counts = Counter(
        str(record["video_first_frame_hash"])
        for record in frozen_records
        if record.get("video_first_frame_hash")
    )
    repeated_frozen_frames = [
        {"first_frame_hash": frame_hash, "view_count": count}
        for frame_hash, count in first_frame_hash_counts.most_common()
        if count > 1
    ]
    all_classification_records = [
        {"classification": name, **record}
        for name, records in classifications.items()
        for record in records
    ]
    classification_path = output_dir / "wrist_static_classifications.json"
    classification_path.write_text(
        json.dumps(all_classification_records, indent=2) + "\n"
    )
    blacklist_path = output_dir / "frozen_wrist_episode_indices.json"
    blacklist_path.write_text(
        json.dumps(frozen_episode_indices, indent=2) + "\n"
    )

    summary = {
        "dataset_root": str(root),
        "total_episodes": len(rows),
        "total_frames": int(sum(int(row["frames"]) for row in rows)),
        "track_integrity_error_count": len(errors),
        "track_integrity_errors": errors[:100],
        "static_background_threshold_px": static_threshold,
        "near_static_wrist_view_count": len(candidates),
        "near_static_episode_count": len(candidate_episode_indices),
        "wrist_video_check_count": len(video_results),
        "wrist_video_check_error_count": len(video_error_keys),
        "wrist_video_frame_count_mismatch_count": len(
            video_frame_count_mismatches
        ),
        "wrist_video_frame_count_mismatches": (
            video_frame_count_mismatches[:100]
        ),
        "source_video_near_static_view_count": len(
            video_near_static_keys
        ),
        "source_video_near_static_episode_count": len(
            {episode for episode, _ in video_near_static_keys}
        ),
        "classification_counts": {
            name: len(records)
            for name, records in classifications.items()
        },
        "classification_episode_counts": {
            name: len(
                {
                    int(record["episode_index"])
                    for record in records
                }
            )
            for name, records in classifications.items()
        },
        "classification_examples": {
            name: records[:30]
            for name, records in classifications.items()
            if records
        },
        "focus_episode_5463": [
            record
            for record in all_classification_records
            if int(record["episode_index"]) == 5463
        ],
        "repeated_frozen_first_frame_groups": repeated_frozen_frames[:30],
        "unique_frozen_first_frames": len(first_frame_hash_counts),
        "frozen_view_breakdown": dict(
            Counter(str(record["view"]) for record in frozen_records)
        ),
        "frozen_episode_breakdown": dict(frozen_episode_breakdown),
        "frozen_episode_frame_count": frozen_frame_count,
        "frozen_episode_frame_fraction": (
            frozen_frame_count
            / sum(int(row["frames"]) for row in rows)
        ),
        "seam_anomaly_count": len(seam_anomalies),
        "seam_anomalies": seam_anomalies[:100],
        "metric_distributions": metric_distributions,
        "episode_metrics_csv": str(csv_path),
        "wrist_video_metrics_csv": (
            str(video_csv_path) if video_results else None
        ),
        "wrist_static_classifications_json": str(classification_path),
        "frozen_wrist_episode_indices_json": str(blacklist_path),
    }
    summary_path = output_dir / "summary.json"
    summary_path.write_text(json.dumps(summary, indent=2) + "\n")
    print(f"Wrote {summary_path}")
    print(f"Wrote {csv_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())