feat(data-only-viz): dataset jsonl+windows
Implements core dataset infrastructure for action-head training: - RawFrame: JSONL frame parsing (ts, session, pid, j3d) - sliding_windows: temporal windows by (session, pid) - DatasetRow: labeled windows with confidence/validation - write/load_dataset_jsonl: JSONL numpy array serialization - split_by_session: stratified train/val/test by session Tests verify load→windows→write→load→split data flows.
This commit is contained in:
L'électron rare
@@ -0,0 +1,94 @@
|
||||
"""Tests for dataset jsonl IO + sliding windows + split."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
|
||||
def _make_session_jsonl(path: Path, n_frames: int = 64) -> None:
|
||||
rng = np.random.default_rng(0)
|
||||
with path.open("w") as f:
|
||||
for t in range(n_frames):
|
||||
row = {"ts": t / 30.0,
|
||||
"session": "sess01",
|
||||
"pid": 1,
|
||||
"j3d": rng.normal(size=(22, 3)).tolist()}
|
||||
f.write(json.dumps(row) + "\n")
|
||||
|
||||
|
||||
def test_load_frames_jsonl(tmp_path: Path) -> None:
|
||||
from data_only_viz.training.dataset import load_frames_jsonl
|
||||
p = tmp_path / "raw.jsonl"
|
||||
_make_session_jsonl(p)
|
||||
frames = load_frames_jsonl(p)
|
||||
assert len(frames) == 64
|
||||
assert frames[0].j3d.shape == (22, 3)
|
||||
assert frames[0].pid == 1
|
||||
assert frames[0].session == "sess01"
|
||||
|
||||
|
||||
def test_sliding_windows(tmp_path: Path) -> None:
|
||||
from data_only_viz.training.dataset import (
|
||||
load_frames_jsonl,
|
||||
sliding_windows,
|
||||
)
|
||||
p = tmp_path / "raw.jsonl"
|
||||
_make_session_jsonl(p, n_frames=64)
|
||||
frames = load_frames_jsonl(p)
|
||||
windows = list(sliding_windows(frames, window_len=16, stride=4))
|
||||
assert len(windows) == 13
|
||||
assert windows[0].j3d_stack.shape == (16, 22, 3)
|
||||
assert windows[0].session == "sess01"
|
||||
|
||||
|
||||
def test_write_and_load_dataset_jsonl(tmp_path: Path) -> None:
|
||||
from data_only_viz.training.dataset import (
|
||||
DatasetRow,
|
||||
load_dataset_jsonl,
|
||||
write_dataset_jsonl,
|
||||
)
|
||||
rng = np.random.default_rng(0)
|
||||
rows = [
|
||||
DatasetRow(
|
||||
window_id=f"sess01_pid1_w{i:04d}",
|
||||
label="debout" if i % 2 == 0 else "danse",
|
||||
j3d_stack=rng.normal(size=(16, 22, 3)).astype(np.float32),
|
||||
session="sess01",
|
||||
pid_local=1,
|
||||
auto_label_confidence=0.8,
|
||||
manually_validated=False,
|
||||
)
|
||||
for i in range(5)
|
||||
]
|
||||
out = tmp_path / "ds.jsonl"
|
||||
write_dataset_jsonl(rows, out)
|
||||
loaded = load_dataset_jsonl(out)
|
||||
assert len(loaded) == 5
|
||||
assert loaded[0].label == "debout"
|
||||
assert loaded[0].j3d_stack.shape == (16, 22, 3)
|
||||
assert np.allclose(loaded[0].j3d_stack, rows[0].j3d_stack, atol=1e-6)
|
||||
|
||||
|
||||
def test_split_by_session(tmp_path: Path) -> None:
|
||||
from data_only_viz.training.dataset import DatasetRow, split_by_session
|
||||
rng = np.random.default_rng(0)
|
||||
rows = []
|
||||
for sess in ("s01", "s02", "s03", "s04", "s05", "s06", "s07"):
|
||||
rows.append(DatasetRow(
|
||||
window_id=f"{sess}_w0", label="debout",
|
||||
j3d_stack=rng.normal(size=(16, 22, 3)).astype(np.float32),
|
||||
session=sess, pid_local=1, auto_label_confidence=0.7,
|
||||
manually_validated=False,
|
||||
))
|
||||
train, val, test = split_by_session(rows, ratios=(0.7, 0.15, 0.15), seed=0)
|
||||
all_sessions = {r.session for r in train + val + test}
|
||||
assert all_sessions == {"s01","s02","s03","s04","s05","s06","s07"}
|
||||
train_s = {r.session for r in train}
|
||||
val_s = {r.session for r in val}
|
||||
test_s = {r.session for r in test}
|
||||
assert train_s.isdisjoint(val_s)
|
||||
assert train_s.isdisjoint(test_s)
|
||||
assert val_s.isdisjoint(test_s)
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Dataset IO + sliding-window extraction + by-session split."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Iterable, Iterator
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RawFrame:
|
||||
ts: float
|
||||
session: str
|
||||
pid: int
|
||||
j3d: np.ndarray # (22, 3) float32
|
||||
|
||||
|
||||
@dataclass
|
||||
class WindowRow:
|
||||
j3d_stack: np.ndarray # (window_len, 22, 3) float32
|
||||
session: str
|
||||
pid_local: int
|
||||
first_ts: float
|
||||
|
||||
|
||||
@dataclass
|
||||
class DatasetRow:
|
||||
window_id: str
|
||||
label: str
|
||||
j3d_stack: np.ndarray # (window_len, 22, 3) float32
|
||||
session: str
|
||||
pid_local: int
|
||||
auto_label_confidence: float
|
||||
manually_validated: bool
|
||||
|
||||
|
||||
def load_frames_jsonl(path: Path) -> list[RawFrame]:
|
||||
rows: list[RawFrame] = []
|
||||
with path.open() as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
d = json.loads(line)
|
||||
rows.append(RawFrame(
|
||||
ts=float(d["ts"]),
|
||||
session=str(d["session"]),
|
||||
pid=int(d["pid"]),
|
||||
j3d=np.asarray(d["j3d"], dtype=np.float32),
|
||||
))
|
||||
return rows
|
||||
|
||||
|
||||
def sliding_windows(frames: list[RawFrame],
|
||||
window_len: int = 16,
|
||||
stride: int = 4) -> Iterator[WindowRow]:
|
||||
"""Yield (session, pid)-grouped windows."""
|
||||
by_key: dict[tuple[str, int], list[RawFrame]] = {}
|
||||
for fr in frames:
|
||||
by_key.setdefault((fr.session, fr.pid), []).append(fr)
|
||||
for (sess, pid), grp in by_key.items():
|
||||
grp.sort(key=lambda r: r.ts)
|
||||
if len(grp) < window_len:
|
||||
continue
|
||||
for start in range(0, len(grp) - window_len + 1, stride):
|
||||
chunk = grp[start:start + window_len]
|
||||
stack = np.stack([c.j3d for c in chunk]).astype(np.float32)
|
||||
yield WindowRow(j3d_stack=stack, session=sess,
|
||||
pid_local=pid, first_ts=chunk[0].ts)
|
||||
|
||||
|
||||
def write_dataset_jsonl(rows: Iterable[DatasetRow], path: Path) -> None:
|
||||
with path.open("w") as f:
|
||||
for r in rows:
|
||||
f.write(json.dumps({
|
||||
"window_id": r.window_id,
|
||||
"label": r.label,
|
||||
"j3d": r.j3d_stack.astype(np.float32).tolist(),
|
||||
"session": r.session,
|
||||
"pid_local": r.pid_local,
|
||||
"auto_label_confidence": float(r.auto_label_confidence),
|
||||
"manually_validated": bool(r.manually_validated),
|
||||
}) + "\n")
|
||||
|
||||
|
||||
def load_dataset_jsonl(path: Path) -> list[DatasetRow]:
|
||||
out: list[DatasetRow] = []
|
||||
with path.open() as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
d = json.loads(line)
|
||||
out.append(DatasetRow(
|
||||
window_id=d["window_id"],
|
||||
label=d["label"],
|
||||
j3d_stack=np.asarray(d["j3d"], dtype=np.float32),
|
||||
session=d["session"],
|
||||
pid_local=int(d["pid_local"]),
|
||||
auto_label_confidence=float(d["auto_label_confidence"]),
|
||||
manually_validated=bool(d["manually_validated"]),
|
||||
))
|
||||
return out
|
||||
|
||||
|
||||
def split_by_session(rows: list[DatasetRow],
|
||||
ratios: tuple[float, float, float] = (0.7, 0.15, 0.15),
|
||||
seed: int = 0,
|
||||
) -> tuple[list[DatasetRow], list[DatasetRow], list[DatasetRow]]:
|
||||
sessions = sorted({r.session for r in rows})
|
||||
rng = random.Random(seed)
|
||||
rng.shuffle(sessions)
|
||||
n = len(sessions)
|
||||
n_train = max(1, int(round(n * ratios[0])))
|
||||
n_val = max(1, int(round(n * ratios[1])))
|
||||
if n_train + n_val >= n:
|
||||
n_val = max(1, n - n_train - 1)
|
||||
train_s = set(sessions[:n_train])
|
||||
val_s = set(sessions[n_train:n_train + n_val])
|
||||
test_s = set(sessions[n_train + n_val:])
|
||||
train = [r for r in rows if r.session in train_s]
|
||||
val = [r for r in rows if r.session in val_s]
|
||||
test = [r for r in rows if r.session in test_s]
|
||||
return train, val, test
|
||||
Reference in New Issue
Block a user