| |
| """ |
| Minimal end-to-end action-recognition pipeline for an Xperience-10M episode. |
| |
| Input: |
| annotation.hdf5 |
| |
| Features: |
| hand joints, body joints, contacts, camera trajectory, IMU summary statistics. |
| |
| Target: |
| caption action_label by default. Use --target subtask for Sub Task labels. |
| |
| Model: |
| Numpy-only multinomial logistic regression. |
| |
| Outputs: |
| metrics.json, per_class_metrics.csv, confusion_matrix.csv, predictions.csv, |
| feature_dataset.npz, model.npz. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import json |
| import math |
| import sys |
| from collections import Counter, OrderedDict |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| workspace_default = Path(__file__).resolve().parents[1] |
| data_default = workspace_default / "data/sample/xperience-10m-sample/annotation.hdf5" |
| out_default = workspace_default / "outputs/min_action_model" |
|
|
| parser = argparse.ArgumentParser(description="Train a minimal action classifier on Ropedia annotation.hdf5.") |
| parser.add_argument("--workspace", type=Path, default=workspace_default, help="Ropedia workspace root.") |
| parser.add_argument("--annotation", type=Path, default=data_default, help="Path to annotation.hdf5.") |
| parser.add_argument("--output-dir", type=Path, default=out_default, help="Output artifact directory.") |
| parser.add_argument("--target", choices=["action", "subtask"], default="action", help="Prediction target.") |
| parser.add_argument("--window-frames", type=int, default=20, help="Frames per training window.") |
| parser.add_argument("--stride-frames", type=int, default=5, help="Stride between windows.") |
| parser.add_argument("--min-label-fraction", type=float, default=0.6, help="Minimum majority-label fraction in a window.") |
| parser.add_argument("--test-fraction", type=float, default=0.25, help="Stratified test fraction.") |
| parser.add_argument("--epochs", type=int, default=800, help="Training epochs.") |
| parser.add_argument("--learning-rate", type=float, default=0.2, help="Softmax learning rate.") |
| parser.add_argument("--l2", type=float, default=1e-3, help="L2 weight decay.") |
| parser.add_argument("--seed", type=int, default=7, help="Random seed.") |
| parser.add_argument("--no-class-weights", action="store_true", help="Disable inverse-frequency class weighting.") |
| return parser.parse_args() |
|
|
|
|
| def add_toolkit_to_path(workspace: Path) -> None: |
| toolkit = workspace / "HOMIE-toolkit" |
| if not toolkit.exists(): |
| raise FileNotFoundError(f"HOMIE-toolkit not found: {toolkit}") |
| sys.path.insert(0, str(toolkit)) |
|
|
|
|
| def portable_path(path: Path, workspace: Path | None = None) -> str: |
| roots = [workspace, Path.cwd()] |
| for root in roots: |
| if root is None: |
| continue |
| try: |
| return path.resolve().relative_to(Path(root).resolve()).as_posix() |
| except (FileNotFoundError, ValueError): |
| continue |
| return path.name |
|
|
|
|
| def temporal_stats(arr: np.ndarray) -> np.ndarray: |
| """Return fixed statistics over time for an array shaped (T, ...).""" |
| arr = np.asarray(arr, dtype=np.float32) |
| if arr.ndim == 0: |
| arr = arr.reshape(1, 1) |
| elif arr.ndim == 1: |
| arr = arr[:, None] |
| flat = arr.reshape(arr.shape[0], -1) |
| flat = np.nan_to_num(flat, nan=0.0, posinf=0.0, neginf=0.0) |
| if flat.shape[0] == 0: |
| raise ValueError("temporal_stats received an empty time axis") |
|
|
| mean = flat.mean(axis=0) |
| std = flat.std(axis=0) |
| amin = flat.min(axis=0) |
| amax = flat.max(axis=0) |
| delta = flat[-1] - flat[0] |
| if flat.shape[0] > 1: |
| vel = np.diff(flat, axis=0) |
| vel_mean = vel.mean(axis=0) |
| vel_std = vel.std(axis=0) |
| else: |
| vel_mean = np.zeros(flat.shape[1], dtype=np.float32) |
| vel_std = np.zeros(flat.shape[1], dtype=np.float32) |
| return np.concatenate([mean, std, amin, amax, delta, vel_mean, vel_std]).astype(np.float32) |
|
|
|
|
| def safe_window(arr: np.ndarray | None, start: int, end: int) -> np.ndarray | None: |
| if arr is None: |
| return None |
| if start >= len(arr): |
| return None |
| return np.asarray(arr[start:min(end, len(arr))]) |
|
|
|
|
| def center_by_body_root(values: np.ndarray, body: np.ndarray | None) -> np.ndarray: |
| if body is None or len(body) != len(values) or body.ndim < 3 or body.shape[-1] != 3: |
| return values |
| root = body[:, :1, :] |
| return values - root |
|
|
|
|
| def extract_window_features(ann: dict, start: int, end: int) -> np.ndarray: |
| body = safe_window(ann.get("smplh_body_joints"), start, end) |
| left = safe_window(ann.get("hand_left_joints"), start, end) |
| right = safe_window(ann.get("hand_right_joints"), start, end) |
| contacts = safe_window(ann.get("contacts"), start, end) |
| cam_t = safe_window(ann.get("t_c2w_all"), start, end) |
|
|
| chunks: list[np.ndarray] = [] |
|
|
| if left is not None: |
| chunks.append(temporal_stats(center_by_body_root(left, body))) |
| if right is not None: |
| chunks.append(temporal_stats(center_by_body_root(right, body))) |
| if body is not None: |
| root = body[:, :1, :] if body.ndim == 3 else 0.0 |
| chunks.append(temporal_stats(body - root)) |
| if contacts is not None: |
| chunks.append(temporal_stats(contacts)) |
| if cam_t is not None: |
| cam_t = cam_t - cam_t[:1] |
| chunks.append(temporal_stats(cam_t)) |
|
|
| imu_accel = ann.get("imu_accel_xyz") |
| imu_gyro = ann.get("imu_gyro_xyz") |
| imu_keyframes = ann.get("imu_keyframe_indices") |
| if imu_accel is not None and imu_gyro is not None and imu_keyframes is not None and len(imu_keyframes) > end - 1: |
| imu_start = int(max(0, imu_keyframes[start])) |
| imu_end = int(min(len(imu_accel), max(imu_start + 1, imu_keyframes[end - 1] + 1))) |
| imu = np.concatenate([imu_accel[imu_start:imu_end], imu_gyro[imu_start:imu_end]], axis=1) |
| chunks.append(temporal_stats(imu)) |
|
|
| if not chunks: |
| raise ValueError("No usable numeric modalities found in annotation.") |
| return np.concatenate(chunks).astype(np.float32) |
|
|
|
|
| def frame_label(info: dict, target: str) -> str: |
| if target == "subtask": |
| label = info.get("theme", "") |
| else: |
| label = info.get("action_label", "") |
| label = str(label).strip() |
| if not label or label.upper() == "N/A": |
| return "" |
| return label |
|
|
|
|
| def majority_label(labels: list[str], min_fraction: float) -> tuple[str, float]: |
| labels = [x for x in labels if x] |
| if not labels: |
| return "", 0.0 |
| label, count = Counter(labels).most_common(1)[0] |
| frac = count / len(labels) |
| if frac < min_fraction: |
| return "", frac |
| return label, frac |
|
|
|
|
| def build_feature_dataset(ann: dict, target: str, window_frames: int, stride_frames: int, min_label_fraction: float): |
| frame_info = ann.get("caption_frame_info_map") |
| if frame_info is None: |
| raise ValueError("No caption_frame_info_map found in annotation.") |
|
|
| n_frames = len(ann["img_names"]) |
| X, y_labels, starts, ends, label_fracs = [], [], [], [], [] |
| for start in range(0, n_frames - window_frames + 1, stride_frames): |
| end = start + window_frames |
| labels = [frame_label(frame_info.get(i, {}), target) for i in range(start, end)] |
| label, frac = majority_label(labels, min_label_fraction) |
| if not label: |
| continue |
| X.append(extract_window_features(ann, start, end)) |
| y_labels.append(label) |
| starts.append(start) |
| ends.append(end - 1) |
| label_fracs.append(frac) |
|
|
| if not X: |
| raise ValueError("No labeled windows were created. Try lowering --min-label-fraction.") |
|
|
| return ( |
| np.stack(X).astype(np.float32), |
| np.asarray(y_labels, dtype=object), |
| np.asarray(starts, dtype=np.int64), |
| np.asarray(ends, dtype=np.int64), |
| np.asarray(label_fracs, dtype=np.float32), |
| ) |
|
|
|
|
| def encode_labels(y_labels: np.ndarray) -> tuple[np.ndarray, list[str]]: |
| seen = OrderedDict() |
| for label in y_labels: |
| if label not in seen: |
| seen[label] = len(seen) |
| class_names = list(seen.keys()) |
| y = np.asarray([seen[label] for label in y_labels], dtype=np.int64) |
| return y, class_names |
|
|
|
|
| def stratified_split(y: np.ndarray, test_fraction: float, seed: int) -> tuple[np.ndarray, np.ndarray]: |
| rng = np.random.default_rng(seed) |
| train_idx, test_idx = [], [] |
| for cls in np.unique(y): |
| idx = np.flatnonzero(y == cls) |
| rng.shuffle(idx) |
| if len(idx) < 2: |
| train_idx.extend(idx.tolist()) |
| continue |
| n_test = int(round(len(idx) * test_fraction)) |
| n_test = max(1, min(n_test, len(idx) - 1)) |
| test_idx.extend(idx[:n_test].tolist()) |
| train_idx.extend(idx[n_test:].tolist()) |
| rng.shuffle(train_idx) |
| rng.shuffle(test_idx) |
| return np.asarray(train_idx, dtype=np.int64), np.asarray(test_idx, dtype=np.int64) |
|
|
|
|
| def fit_scaler(X: np.ndarray) -> tuple[np.ndarray, np.ndarray]: |
| mean = X.mean(axis=0) |
| std = X.std(axis=0) |
| std = np.where(std < 1e-6, 1.0, std) |
| return mean.astype(np.float32), std.astype(np.float32) |
|
|
|
|
| def softmax(logits: np.ndarray) -> np.ndarray: |
| logits = logits - logits.max(axis=1, keepdims=True) |
| exp = np.exp(logits) |
| return exp / exp.sum(axis=1, keepdims=True) |
|
|
|
|
| def train_softmax_classifier( |
| X: np.ndarray, |
| y: np.ndarray, |
| n_classes: int, |
| epochs: int, |
| lr: float, |
| l2: float, |
| use_class_weights: bool, |
| seed: int, |
| ) -> tuple[np.ndarray, np.ndarray, list[dict]]: |
| rng = np.random.default_rng(seed) |
| n, d = X.shape |
| W = rng.normal(0.0, 0.01, size=(d, n_classes)).astype(np.float32) |
| b = np.zeros(n_classes, dtype=np.float32) |
| onehot = np.eye(n_classes, dtype=np.float32)[y] |
|
|
| if use_class_weights: |
| counts = np.bincount(y, minlength=n_classes).astype(np.float32) |
| weights_by_class = n / np.maximum(counts, 1.0) / n_classes |
| sample_weights = weights_by_class[y] |
| else: |
| sample_weights = np.ones(n, dtype=np.float32) |
| sample_weights = sample_weights / sample_weights.mean() |
|
|
| history = [] |
| report_every = max(1, epochs // 10) |
| for epoch in range(1, epochs + 1): |
| logits = X @ W + b |
| probs = softmax(logits) |
| weighted_diff = (probs - onehot) * sample_weights[:, None] / n |
| grad_W = X.T @ weighted_diff + l2 * W |
| grad_b = weighted_diff.sum(axis=0) |
| W -= lr * grad_W |
| b -= lr * grad_b |
|
|
| if epoch == 1 or epoch == epochs or epoch % report_every == 0: |
| p_true = np.clip(probs[np.arange(n), y], 1e-9, 1.0) |
| loss = float(-(sample_weights * np.log(p_true)).mean() + 0.5 * l2 * float(np.sum(W * W))) |
| acc = float(np.mean(np.argmax(probs, axis=1) == y)) |
| history.append({"epoch": epoch, "loss": loss, "train_accuracy": acc}) |
| return W.astype(np.float32), b.astype(np.float32), history |
|
|
|
|
| def predict(X: np.ndarray, W: np.ndarray, b: np.ndarray) -> tuple[np.ndarray, np.ndarray]: |
| probs = softmax(X @ W + b) |
| return np.argmax(probs, axis=1), probs |
|
|
|
|
| def compute_metrics(y_true: np.ndarray, y_pred: np.ndarray, class_names: list[str]) -> tuple[dict, list[dict], np.ndarray]: |
| n_classes = len(class_names) |
| cm = np.zeros((n_classes, n_classes), dtype=np.int64) |
| for t, p in zip(y_true, y_pred): |
| cm[int(t), int(p)] += 1 |
|
|
| rows = [] |
| recalls, f1s, weighted_f1_total = [], [], 0.0 |
| support_total = int(cm.sum()) |
| for i, name in enumerate(class_names): |
| tp = int(cm[i, i]) |
| support = int(cm[i, :].sum()) |
| pred_count = int(cm[:, i].sum()) |
| precision = tp / pred_count if pred_count else 0.0 |
| recall = tp / support if support else 0.0 |
| f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0 |
| if support: |
| recalls.append(recall) |
| f1s.append(f1) |
| weighted_f1_total += f1 * support |
| rows.append({ |
| "class_id": i, |
| "class_name": name, |
| "support": support, |
| "predicted": pred_count, |
| "precision": precision, |
| "recall": recall, |
| "f1": f1, |
| }) |
|
|
| accuracy = float(np.mean(y_true == y_pred)) if len(y_true) else 0.0 |
| macro_f1 = float(np.mean(f1s)) if f1s else 0.0 |
| balanced_accuracy = float(np.mean(recalls)) if recalls else 0.0 |
| weighted_f1 = float(weighted_f1_total / support_total) if support_total else 0.0 |
| metrics = { |
| "accuracy": accuracy, |
| "balanced_accuracy": balanced_accuracy, |
| "macro_f1": macro_f1, |
| "weighted_f1": weighted_f1, |
| "num_eval_windows": int(len(y_true)), |
| "num_classes": n_classes, |
| } |
| return metrics, rows, cm |
|
|
|
|
| def write_csv(path: Path, rows: list[dict], fieldnames: list[str]) -> None: |
| with path.open("w", newline="", encoding="utf-8") as fp: |
| writer = csv.DictWriter(fp, fieldnames=fieldnames, lineterminator="\n") |
| writer.writeheader() |
| writer.writerows(rows) |
|
|
|
|
| def save_artifacts( |
| output_dir: Path, |
| X: np.ndarray, |
| y: np.ndarray, |
| y_labels: np.ndarray, |
| starts: np.ndarray, |
| ends: np.ndarray, |
| label_fracs: np.ndarray, |
| train_idx: np.ndarray, |
| test_idx: np.ndarray, |
| class_names: list[str], |
| mean: np.ndarray, |
| std: np.ndarray, |
| W: np.ndarray, |
| b: np.ndarray, |
| history: list[dict], |
| metrics: dict, |
| per_class_rows: list[dict], |
| cm: np.ndarray, |
| y_pred: np.ndarray, |
| probs: np.ndarray, |
| args: argparse.Namespace, |
| ) -> None: |
| output_dir.mkdir(parents=True, exist_ok=True) |
|
|
| np.savez_compressed( |
| output_dir / "feature_dataset.npz", |
| X=X, |
| y=y, |
| labels=y_labels.astype(str), |
| start_frame=starts, |
| end_frame=ends, |
| label_fraction=label_fracs, |
| train_idx=train_idx, |
| test_idx=test_idx, |
| class_names=np.asarray(class_names, dtype=object), |
| ) |
| np.savez_compressed(output_dir / "model.npz", mean=mean, std=std, W=W, b=b, class_names=np.asarray(class_names, dtype=object)) |
|
|
| metadata = { |
| "annotation": portable_path(args.annotation, args.workspace), |
| "target": args.target, |
| "window_frames": args.window_frames, |
| "stride_frames": args.stride_frames, |
| "min_label_fraction": args.min_label_fraction, |
| "test_fraction": args.test_fraction, |
| "epochs": args.epochs, |
| "learning_rate": args.learning_rate, |
| "l2": args.l2, |
| "class_weights": not args.no_class_weights, |
| "num_windows": int(len(y)), |
| "num_features": int(X.shape[1]), |
| "num_train_windows": int(len(train_idx)), |
| "num_test_windows": int(len(test_idx)), |
| "classes": class_names, |
| "history": history, |
| } |
| (output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") |
| (output_dir / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8") |
|
|
| write_csv( |
| output_dir / "per_class_metrics.csv", |
| per_class_rows, |
| ["class_id", "class_name", "support", "predicted", "precision", "recall", "f1"], |
| ) |
|
|
| with (output_dir / "confusion_matrix.csv").open("w", newline="", encoding="utf-8") as fp: |
| writer = csv.writer(fp, lineterminator="\n") |
| writer.writerow(["true\\pred"] + class_names) |
| for i, name in enumerate(class_names): |
| writer.writerow([name] + [int(v) for v in cm[i]]) |
|
|
| pred_rows = [] |
| pred_lookup = {int(idx): k for k, idx in enumerate(test_idx)} |
| for idx in test_idx: |
| idx = int(idx) |
| k = pred_lookup[idx] |
| pred_id = int(y_pred[k]) |
| true_id = int(y[idx]) |
| pred_rows.append({ |
| "window_index": idx, |
| "start_frame": int(starts[idx]), |
| "end_frame": int(ends[idx]), |
| "true_label": class_names[true_id], |
| "predicted_label": class_names[pred_id], |
| "confidence": float(probs[k, pred_id]), |
| "correct": int(pred_id == true_id), |
| "label_fraction": float(label_fracs[idx]), |
| }) |
| write_csv( |
| output_dir / "predictions.csv", |
| pred_rows, |
| ["window_index", "start_frame", "end_frame", "true_label", "predicted_label", "confidence", "correct", "label_fraction"], |
| ) |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| add_toolkit_to_path(args.workspace) |
| from data_loader import load_from_annotation_hdf5 |
|
|
| if not args.annotation.exists(): |
| raise FileNotFoundError(f"annotation.hdf5 not found: {args.annotation}") |
|
|
| print(f"Loading annotation: {args.annotation}") |
| ann = load_from_annotation_hdf5(args.annotation, 0, None, load_slam_point_cloud=False) |
|
|
| print("Building windowed feature dataset") |
| X, y_labels, starts, ends, label_fracs = build_feature_dataset( |
| ann, |
| target=args.target, |
| window_frames=args.window_frames, |
| stride_frames=args.stride_frames, |
| min_label_fraction=args.min_label_fraction, |
| ) |
| y, class_names = encode_labels(y_labels) |
| train_idx, test_idx = stratified_split(y, args.test_fraction, args.seed) |
| if len(test_idx) == 0: |
| raise ValueError("No test windows available. Lower --test-fraction or use more data.") |
|
|
| mean, std = fit_scaler(X[train_idx]) |
| X_scaled = (X - mean) / std |
|
|
| print(f"Windows: {len(y)} total, {len(train_idx)} train, {len(test_idx)} test") |
| print(f"Features: {X.shape[1]}, classes: {len(class_names)}") |
| for name, count in Counter(y_labels).most_common(): |
| print(f" {count:4d} windows {name}") |
|
|
| print("Training softmax classifier") |
| W, b, history = train_softmax_classifier( |
| X_scaled[train_idx], |
| y[train_idx], |
| n_classes=len(class_names), |
| epochs=args.epochs, |
| lr=args.learning_rate, |
| l2=args.l2, |
| use_class_weights=not args.no_class_weights, |
| seed=args.seed, |
| ) |
|
|
| y_pred, probs = predict(X_scaled[test_idx], W, b) |
| metrics, per_class_rows, cm = compute_metrics(y[test_idx], y_pred, class_names) |
|
|
| majority_class = Counter(y[train_idx]).most_common(1)[0][0] |
| metrics["majority_baseline_accuracy"] = float(np.mean(y[test_idx] == majority_class)) |
| metrics["train_final_accuracy"] = history[-1]["train_accuracy"] if history else math.nan |
| metrics["train_final_loss"] = history[-1]["loss"] if history else math.nan |
|
|
| save_artifacts( |
| args.output_dir, |
| X, |
| y, |
| y_labels, |
| starts, |
| ends, |
| label_fracs, |
| train_idx, |
| test_idx, |
| class_names, |
| mean, |
| std, |
| W, |
| b, |
| history, |
| metrics, |
| per_class_rows, |
| cm, |
| y_pred, |
| probs, |
| args, |
| ) |
|
|
| print("\nEvaluation") |
| print(f" accuracy: {metrics['accuracy']:.4f}") |
| print(f" balanced_accuracy: {metrics['balanced_accuracy']:.4f}") |
| print(f" macro_f1: {metrics['macro_f1']:.4f}") |
| print(f" weighted_f1: {metrics['weighted_f1']:.4f}") |
| print(f" majority_baseline: {metrics['majority_baseline_accuracy']:.4f}") |
| print(f"\nArtifacts written to: {args.output_dir}") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
