TALOS-NIO/cache_builder.py at main · Hinedes/TALOS-NIO · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
"""
cache_builder.py -- One-time VRS -> NPZ cache for TALOS NIO
Saves pre-windowing arrays (aligned IMU + GT) to golden/cache/<seq_id>.npz
Run once. Then incremental_train.py reads from cache, VRS never touched again.
"""
import json
import numpy as np
from pathlib import Path
from projectaria_tools.core import data_provider
from nymeria_loader import (
    load_imu_stream, align_imu_streams, load_gt_trajectory,
    interpolate_gt, SID_RIGHT, SID_LEFT, TARGET_HZ
)

RAW_ROOT_CANDIDATES = [
    Path('/home/iclab/TALOS/nymeria'),
    Path('/mnt/c/TALOS/nymeria'),
]
CACHE_DIR = Path('/home/iclab/TALOS/golden/cache')
MANIFEST_CANDIDATES = [
    Path('/home/iclab/TALOS/Nymeria_download_urls.json'),
    Path('/mnt/c/TALOS/Nymeria_download_urls.json'),
]


def first_existing(paths):
    for path in paths:
        if path.exists():
            return path
    return paths[0]


ROOT = first_existing(RAW_ROOT_CANDIDATES)
MANIFEST = first_existing(MANIFEST_CANDIDATES)

CACHE_DIR.mkdir(parents=True, exist_ok=True)

if not ROOT.exists():
    raise FileNotFoundError(f"Raw Nymeria root not found. Checked: {RAW_ROOT_CANDIDATES}")

if not MANIFEST.exists():
    raise FileNotFoundError(f"Nymeria manifest not found. Checked: {MANIFEST_CANDIDATES}")

manifest = json.loads(MANIFEST.read_text())

print(f"[CACHE ROOT] {ROOT}")
print(f"[MANIFEST]    {MANIFEST}")

for seq_id, entry in manifest["sequences"].items():
    out_path = CACHE_DIR / f"Nymeria_v0.0_{seq_id}_recording_head.npz"
    if out_path.exists():
        print(f"[SKIP] {seq_id[:40]}")
        continue

    seq_path = ROOT / f"Nymeria_v0.0_{seq_id}_recording_head" / "recording_head"
    vrs_path  = seq_path / 'data' / 'motion.vrs'
    traj_path = seq_path / 'mps' / 'slam' / 'closed_loop_trajectory.csv'

    if not vrs_path.exists():
        print(f"[MISSING] {seq_id[:40]} -- skipping")
        continue

    print(f"[CACHE] {seq_id[:40]}...")
    try:
        dp = data_provider.create_vrs_data_provider(str(vrs_path))
        device_calib = dp.get_device_calibration()

        T_r = device_calib.get_transform_device_sensor("imu-right")
        R_r = T_r.rotation().to_matrix().astype(np.float32)
        T_l = device_calib.get_transform_device_sensor("imu-left")
        R_l = T_l.rotation().to_matrix().astype(np.float32)

        ts_right, imu_right = load_imu_stream(dp, SID_RIGHT, R_r)
        ts_left,  imu_left  = load_imu_stream(dp, SID_LEFT,  R_l)

        grid_ns, imu1_reg, imu2_reg = align_imu_streams(
            ts_right, imu_right, ts_left, imu_left, TARGET_HZ)

        gt_ts, gt_pos, gt_quat = load_gt_trajectory(traj_path)
        pos_at_imu, quat_at_imu = interpolate_gt(gt_ts, gt_pos, gt_quat, grid_ns)

        mask        = (grid_ns >= gt_ts[0]) & (grid_ns <= gt_ts[-1])
        grid_ns     = grid_ns[mask]
        imu1_reg    = imu1_reg[mask]
        imu2_reg    = imu2_reg[mask]
        pos_at_imu  = pos_at_imu[mask]
        quat_at_imu = quat_at_imu[mask]

        np.savez_compressed(out_path,
            grid_ns=grid_ns,
            imu1=imu1_reg,
            imu2=imu2_reg,
            pos=pos_at_imu,
            quat=quat_at_imu,
        )
        duration_s = (grid_ns[-1] - grid_ns[0]) / 1e9
        print(f"  -> {len(grid_ns)} samples, {duration_s:.1f}s -- saved.")
    except Exception as e:
        print(f"  !! FAILED: {e}")

print("\nDone. Run incremental_train.py -- VRS never opens again.")