File size: 4,115 Bytes
1b0d082 | 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 | #!/usr/bin/env python3
"""Load the challenge splits into memory as (N, 200, 6) float32 windows.
Ported from sample_code/ts-tcmpio-unified-temporal-resnet.ipynb (cells 1-2, 11),
with the Kaggle-mounted paths replaced by the local data root. The whole train
split is 82k windows = 0.39 GB, so everything is held in RAM (and later on the
GPU) rather than streamed.
"""
from __future__ import annotations
import os
from pathlib import Path
import numpy as np
import pandas as pd
WINDOW_SIZE = 200
PLATFORM_TO_ID = {"car": 0, "dog": 1, "drone": 2, "human": 3}
PLATFORMS = tuple(PLATFORM_TO_ID)
# Kaggle mounts the competition at /kaggle/input/<slug>; set TARTANIMU_DATA there
DATA_ROOT = Path(os.environ.get("TARTANIMU_DATA", "data/tartan-imu-challenge-iros2026"))
CACHE_SCHEMA = 2 # 2: DenseWindows carries the gravity-direction stream
def _read(path: Path) -> pd.DataFrame:
return pd.read_csv(path).dropna(axis=1, how="all")
def load_index(data_root: Path = DATA_ROOT) -> dict:
idx = data_root / "index"
train = _read(idx / "train_windows.csv").merge(_read(idx / "train_targets.csv"), on="window_id")
val = _read(idx / "val_windows.csv").merge(_read(idx / "val_targets.csv"), on="window_id")
test = _read(idx / "test_windows.csv")
train["platform_id"] = train["platform"].map(PLATFORM_TO_ID)
val["platform_id"] = val["platform"].map(PLATFORM_TO_ID)
paths = {split: {p.stem: str(p) for p in (data_root / split).rglob("*.npz")}
for split in ("train", "val", "test")}
out = {}
for split, frame in (("train", train), ("val", val), ("test", test)):
frame = frame.reset_index(drop=True)
frame["traj_id"] = frame["traj_id"].astype(str).str.replace(".npz", "", regex=False)
frame["file_path"] = frame["traj_id"].map(paths[split])
if frame["file_path"].isna().any():
raise SystemExit(f"{split}: {frame['file_path'].isna().sum()} windows have no .npz")
out[split] = frame
return out
def cache_path(frame: pd.DataFrame, kind: str) -> Path | None:
"""Where to memoise a decoded split, or None when caching is off.
Keyed by the frame's own contents, so no call site has to name its split and
a mismatched cache cannot be picked up. Set TARTANIMU_CACHE to a writable,
node-local directory: on Kaggle the competition mounts over the network with
no page cache, and every run re-decoding 395 .npz files dominates the job.
"""
root = os.environ.get("TARTANIMU_CACHE")
if not root:
return None
Path(root).mkdir(parents=True, exist_ok=True)
# CACHE_SCHEMA must change whenever the cached arrays gain or lose a field,
# or a stale cache is silently reused with the wrong contents
key = (f"{kind}_v{CACHE_SCHEMA}_{len(frame)}"
f"_{frame['window_id'].iloc[0]}_{frame['window_id'].iloc[-1]}")
return Path(root) / f"{key}.npy"
def build_windows(frame: pd.DataFrame) -> np.ndarray:
"""(len(frame), 200, 6) float32, row i matching frame.iloc[i]."""
cache = cache_path(frame, "win")
if cache is not None and cache.exists():
return np.load(cache, mmap_mode="r")
X = np.empty((len(frame), WINDOW_SIZE, 6), dtype=np.float32)
for file_path, rows in frame.groupby("file_path", sort=False):
with np.load(file_path) as d:
imu = d["imu"]
starts = rows["win_idx"].to_numpy() * WINDOW_SIZE
X[rows.index.to_numpy()] = np.stack([imu[s:s + WINDOW_SIZE] for s in starts])
if cache is not None:
np.save(cache, X)
return X
def sample_weights(train: pd.DataFrame) -> np.ndarray:
"""Notebook's WeightedRandomSampler weights.
1 / (windows in this trajectory * trajectories in this platform) makes every
platform equally likely and every trajectory within a platform equally
likely -- which is exactly how the competition macro-averages.
"""
traj_windows = train.groupby("traj_id")["window_id"].transform("size")
plat_trajs = train.groupby("platform")["traj_id"].transform("nunique")
return (1.0 / (traj_windows * plat_trajs)).to_numpy(np.float64)
|