| |
| """Download one validated LRF-IMU fold pair and generate standardized IMU windows.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import torch |
| from huggingface_hub import hf_hub_download |
|
|
| from lrf_imu.checkpoints import load_flow_checkpoint, load_vae_checkpoint |
| from lrf_imu.generation import sample_reverse_euler |
|
|
|
|
| DEFAULT_REPO_ID = "Aminrezaei/LRF-IMU" |
| ACTIVITIES = { |
| 0: {"name": "walking", "realdisp_code": 1}, |
| 1: {"name": "running", "realdisp_code": 3}, |
| 2: {"name": "jump_up", "realdisp_code": 4}, |
| 3: {"name": "cycling", "realdisp_code": 33}, |
| } |
| ACTIVITY_IDS = {details["name"]: class_id for class_id, details in ACTIVITIES.items()} |
| SENSOR_CHANNELS = {"six_channel": 6, "accelerometer_only": 3} |
| CONFIG_NAMES = { |
| "six_channel": "six_channel_160_40.yaml", |
| "accelerometer_only": "accelerometer_only_160_40.yaml", |
| } |
|
|
|
|
| def sha256_file(path: Path) -> str: |
| digest = hashlib.sha256() |
| with path.open("rb") as stream: |
| for block in iter(lambda: stream.read(1024 * 1024), b""): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def download_file( |
| *, |
| repo_id: str, |
| filename: str, |
| revision: str, |
| cache_dir: Path | None, |
| ) -> Path: |
| return Path( |
| hf_hub_download( |
| repo_id=repo_id, |
| filename=filename, |
| revision=revision, |
| cache_dir=None if cache_dir is None else str(cache_dir), |
| ) |
| ) |
|
|
|
|
| def read_json(path: Path) -> dict[str, Any]: |
| with path.open("r", encoding="utf-8") as stream: |
| value = json.load(stream) |
| if not isinstance(value, dict): |
| raise ValueError(f"expected a JSON object in {path.name}") |
| return value |
|
|
|
|
| def require_hash(path: Path, expected: str, label: str) -> str: |
| actual = sha256_file(path) |
| if actual.lower() != str(expected).lower(): |
| raise ValueError(f"{label} SHA-256 mismatch: expected {expected}, got {actual}") |
| return actual |
|
|
|
|
| def resolve_device(requested: str) -> torch.device: |
| if requested == "auto": |
| return torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| if requested == "cuda" and not torch.cuda.is_available(): |
| raise RuntimeError("CUDA was requested but torch.cuda.is_available() is false") |
| return torch.device(requested) |
|
|
|
|
| def selected_model(index: dict[str, Any], sensor: str, fold: int) -> dict[str, Any]: |
| for entry in index.get("models", []): |
| if ( |
| entry.get("sensor_configuration") == sensor |
| and int(entry.get("fold", -1)) == fold |
| ): |
| return entry |
| raise ValueError(f"no checkpoint pair for sensor={sensor!r}, fold={fold:02d}") |
|
|
|
|
| def build_parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser( |
| description=( |
| "Download one matched LRF-IMU fold pair, verify its hashes, and " |
| "generate standardized VAE-decoder-space IMU windows." |
| ) |
| ) |
| parser.add_argument("--repo-id", default=DEFAULT_REPO_ID) |
| parser.add_argument("--revision", default="main") |
| parser.add_argument( |
| "--sensor", |
| choices=tuple(SENSOR_CHANNELS), |
| default="six_channel", |
| ) |
| parser.add_argument( |
| "--fold", |
| type=int, |
| choices=(1, 2, 3, 5, 8, 9, 10, 11, 12, 13, 14, 16), |
| default=1, |
| help="held-out REALDISP participant/fold", |
| ) |
| parser.add_argument( |
| "--activity", |
| choices=tuple(ACTIVITY_IDS), |
| default="walking", |
| ) |
| parser.add_argument("--count", type=int, default=4) |
| parser.add_argument("--steps", type=int, default=10) |
| parser.add_argument("--seed", type=int, default=42) |
| parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto") |
| parser.add_argument("--cache-dir", type=Path) |
| parser.add_argument("--output", type=Path, default=Path("synthetic_imu.npz")) |
| parser.add_argument("--overwrite", action="store_true") |
| return parser |
|
|
|
|
| def main() -> int: |
| args = build_parser().parse_args() |
| class_id = ACTIVITY_IDS[args.activity] |
| if args.count < 1: |
| raise ValueError("--count must be positive") |
| if args.steps < 1: |
| raise ValueError("--steps must be positive") |
|
|
| output = args.output.expanduser().resolve() |
| metadata_output = output.with_suffix(".metadata.json") |
| if not args.overwrite: |
| existing = [path for path in (output, metadata_output) if path.exists()] |
| if existing: |
| raise FileExistsError( |
| "refusing to overwrite existing output: " |
| + ", ".join(str(path) for path in existing) |
| ) |
|
|
| cache_dir = ( |
| None if args.cache_dir is None else args.cache_dir.expanduser().resolve() |
| ) |
| index_path = download_file( |
| repo_id=args.repo_id, |
| filename="model_index.json", |
| revision=args.revision, |
| cache_dir=cache_dir, |
| ) |
| index = read_json(index_path) |
| model = selected_model(index, args.sensor, args.fold) |
|
|
| config_filename = str(model["config"]) |
| expected_config = f"configs/{CONFIG_NAMES[args.sensor]}" |
| if config_filename != expected_config: |
| raise ValueError( |
| f"index config mismatch: expected {expected_config}, got {config_filename}" |
| ) |
|
|
| selected_paths = { |
| "config": config_filename, |
| "manifest": str(model["manifest"]), |
| "vae": str(model["vae"]["path"]), |
| "flow": str(model["flow"]["path"]), |
| } |
| downloaded = { |
| key: download_file( |
| repo_id=args.repo_id, |
| filename=filename, |
| revision=args.revision, |
| cache_dir=cache_dir, |
| ) |
| for key, filename in selected_paths.items() |
| } |
|
|
| manifest = read_json(downloaded["manifest"]) |
| if manifest.get("sensor_configuration") != args.sensor: |
| raise ValueError("manifest sensor does not match the requested sensor") |
| if int(manifest.get("fold", -1)) != args.fold: |
| raise ValueError("manifest fold does not match the requested fold") |
|
|
| index_config = index["configs"][args.sensor] |
| require_hash(downloaded["config"], index_config["sha256"], "config") |
| vae_hash = require_hash( |
| downloaded["vae"], manifest["files"]["vae"]["sha256"], "VAE checkpoint" |
| ) |
| flow_hash = require_hash( |
| downloaded["flow"], manifest["files"]["flow"]["sha256"], "Flow checkpoint" |
| ) |
| if vae_hash != str(model["vae"]["sha256"]).lower(): |
| raise ValueError("VAE hash differs between model_index.json and manifest.json") |
| if flow_hash != str(model["flow"]["sha256"]).lower(): |
| raise ValueError("Flow hash differs between model_index.json and manifest.json") |
|
|
| device = resolve_device(args.device) |
| channels = SENSOR_CHANNELS[args.sensor] |
| vae, vae_info = load_vae_checkpoint( |
| downloaded["vae"], |
| channels=channels, |
| latent_channels=48, |
| down_levels=2, |
| device=device, |
| ) |
| flow, flow_info = load_flow_checkpoint( |
| downloaded["flow"], |
| channels=channels, |
| latent_channels=48, |
| num_classes=4, |
| model_width=256, |
| device=device, |
| ) |
| vae.eval() |
| flow.eval() |
|
|
| if vae_info.latent_time_steps != 40 or flow_info.model_width != 256: |
| raise ValueError("loaded checkpoint geometry differs from the published index") |
| labels = torch.full( |
| (args.count,), |
| class_id, |
| dtype=torch.long, |
| device=device, |
| ) |
| with torch.inference_mode(): |
| latents = sample_reverse_euler( |
| flow, |
| labels, |
| (48, 40), |
| num_steps=args.steps, |
| seed=args.seed, |
| device=device, |
| ) |
| samples = vae.decode(latents).detach().cpu().numpy().astype(np.float32) |
| label_values = labels.detach().cpu().numpy().astype(np.int64) |
|
|
| expected_shape = (args.count, channels, 160) |
| if samples.shape != expected_shape: |
| raise ValueError( |
| f"decoded shape mismatch: expected {expected_shape}, got {samples.shape}" |
| ) |
|
|
| output.parent.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed(output, samples=samples, labels=label_values) |
| output_hash = sha256_file(output) |
| activity = ACTIVITIES[class_id] |
| metadata = { |
| "schema_version": "1.0", |
| "model_id": args.repo_id, |
| "revision": args.revision, |
| "sensor_configuration": args.sensor, |
| "fold": args.fold, |
| "held_out_subject": args.fold, |
| "activity": { |
| "class_id": class_id, |
| "name": activity["name"], |
| "realdisp_code": activity["realdisp_code"], |
| }, |
| "generation": { |
| "solver": "explicit_reverse_euler", |
| "steps": args.steps, |
| "paper_default_steps": 10, |
| "uses_paper_step_count": args.steps == 10, |
| "seed": args.seed, |
| "device": str(device), |
| }, |
| "signal": { |
| "shape": list(samples.shape), |
| "dtype": str(samples.dtype), |
| "orientation": "B,C,T", |
| "sampling_frequency_hz": 50, |
| "window_samples": 160, |
| "coordinate_system": "standardized_vae_decoder_space", |
| "physical_units_assigned": False, |
| }, |
| "checkpoints": { |
| "vae": { |
| "filename": selected_paths["vae"], |
| "sha256": vae_hash, |
| }, |
| "flow": { |
| "filename": selected_paths["flow"], |
| "sha256": flow_hash, |
| }, |
| }, |
| "output": { |
| "filename": output.name, |
| "sha256": output_hash, |
| }, |
| "exact_paper_reproduction": False, |
| } |
| metadata_output.write_text( |
| json.dumps(metadata, indent=2, ensure_ascii=False) + "\n", |
| encoding="utf-8", |
| ) |
| print(f"Wrote {output}") |
| print(f"Wrote {metadata_output}") |
| print( |
| "Signal values are standardized VAE-decoder outputs; " |
| "no physical units are assigned." |
| ) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
