#!/usr/bin/env python3 """ PHASE 4 -- reusable preprocessing pipeline. VIDEO / KEYPOINTS -> landmark packing (75 landmarks x 3 coords) -> geometry units fix (measured; see backend/preprocessing/geometry.py) -> per-frame normalisation (origin = mid-shoulder, scale = shoulder width) -> fixed-length sequence (SEQUENCE_LENGTH frames, linear resample) -> optional landmark mask (config.MASK_ENABLED) -> .npz tensors + class mapping This module is the *only* place that turns dataset rows into model input, so the three models provably receive identical data. It is also used by ``training/robustness.py`` and by the live pipeline (through ``backend.preprocessing.landmarks`` / ``backend.preprocessing.sequence``), which is why the geometry and normalisation code lives in ``backend/preprocessing/`` and not here. Outputs (datasets/processed/) ---------------------------- sequences_.npz X_train/X_val/X_test + y_* + clip metadata preprocess_summary.json sizes, timings, skipped clips, label mapping hash Usage ----- python training/preprocess.py # primary protocol python training/preprocess.py --protocol official # comparability split python training/preprocess.py --augment-copies 2 # extra train copies python training/preprocess.py --limit 50 # smoke test """ from __future__ import annotations import argparse import hashlib import json import logging import sys import time import traceback from pathlib import Path from typing import Dict, List, Optional, Sequence, Tuple import numpy as np sys.path.insert(0, str(Path(__file__).resolve().parents[1])) import config # noqa: E402 from backend.preprocessing import include50, landmarks as lm, sequence as seq_ops # noqa: E402 LOG = logging.getLogger("preprocess") # -------------------------------------------------------------------------- # Core # -------------------------------------------------------------------------- def clip_to_sequence(coords: np.ndarray, mask: np.ndarray, length: Optional[int] = None) -> Optional[np.ndarray]: """(T,75,3) + (T,75) -> (length, 225) normalised feature sequence. Returns ``None`` when the clip carries fewer than ``config.MIN_VALID_FRAMES`` frames that contain any landmark at all (such clips are unusable and are counted in the summary rather than silently zero-padded). """ length = int(length or config.SEQUENCE_LENGTH) if coords.ndim != 3 or coords.shape[0] == 0: return None usable = (np.asarray(mask) > 0).any(axis=1) if int(usable.sum()) < config.MIN_VALID_FRAMES: return None feats, _ = lm.normalise_sequence(coords.astype(np.float32), np.asarray(mask, dtype=np.float32)) feats = seq_ops.resample_sequence(feats, length, method=config.SEQ_RESAMPLE_METHOD) if feats.shape != (length, config.MODEL_INPUT_DIM if config.MASK_ENABLED else config.FEATURE_DIM): return None return feats.astype(np.float32) def build_arrays(index, split: str, augment_copies: int = 0, seed: int = config.RANDOM_SEED, limit: Optional[int] = None) -> Tuple[np.ndarray, np.ndarray, List[Dict], List[Dict]]: """Build (X, y) arrays for one split plus per-clip metadata. Augmentation is applied *offline* to the training split only, with a seeded RNG, so that every model sees byte-identical inputs. Validation and test are never augmented. """ rows = index[index["split"] == split].to_dict("records") if limit: rows = rows[:limit] classes = {c: i for i, c in enumerate(include50.load_vocabulary())} X: List[np.ndarray] = [] y: List[int] = [] meta: List[Dict] = [] skipped: List[Dict] = [] rng = np.random.default_rng(seed + {"train": 0, "val": 1, "test": 2}.get(split, 3)) for row in rows: try: coords, mask, info = include50.read_keypoints_for(row["FilePath"]) except Exception as exc: skipped.append({"filepath": row["FilePath"], "reason": f"{type(exc).__name__}: {exc}"}) continue seq = clip_to_sequence(coords, mask) if seq is None: skipped.append({"filepath": row["FilePath"], "reason": "too few usable frames"}) continue label = classes[row["label"]] copies = [seq] + [seq_ops.augment(seq, rng) for _ in range(augment_copies)] if split == "train" else [seq] for k, s in enumerate(copies): X.append(s) y.append(label) meta.append({ "filepath": row["FilePath"], "label": row["label"], "class_index": label, "split": split, "copy": k, "frames_source": int(coords.shape[0]), "x_axis_scale": float(info.get("x_axis_scale", 1.0)), }) if not X: return (np.zeros((0, config.SEQUENCE_LENGTH, config.FEATURE_DIM), dtype=np.float32), np.zeros((0,), dtype=np.int64), meta, skipped) return np.stack(X).astype(np.float32), np.array(y, dtype=np.int64), meta, skipped def label_mapping() -> Dict[str, object]: vocab = include50.load_vocabulary() payload = {"classes": vocab, "index": {c: i for i, c in enumerate(vocab)}, "num_classes": len(vocab)} payload["sha1"] = hashlib.sha1("|".join(vocab).encode()).hexdigest() return payload def run(protocol: str, augment_copies: int, seed: int, limit: Optional[int], out_dir: Optional[Path] = None) -> Path: out_dir = Path(out_dir or config.PROCESSED_DIR) out_dir.mkdir(parents=True, exist_ok=True) index, split_summary = include50.build_index(protocol=protocol) mapping = label_mapping() LOG.info("protocol=%s train=%d val=%d test=%d classes=%d", protocol, split_summary.train, split_summary.val, split_summary.test, mapping["num_classes"]) t0 = time.time() arrays: Dict[str, np.ndarray] = {} all_meta: Dict[str, List[Dict]] = {} all_skipped: Dict[str, List[Dict]] = {} for split, copies in (("train", augment_copies), ("val", 0), ("test", 0)): X, y, meta, skipped = build_arrays(index, split, augment_copies=copies, seed=seed, limit=limit) arrays[f"X_{split}"] = X arrays[f"y_{split}"] = y all_meta[split] = meta all_skipped[split] = skipped LOG.info(" %-5s X=%s y=%s (skipped %d)", split, X.shape, y.shape, len(skipped)) npz_path = out_dir / f"sequences_{protocol}.npz" np.savez_compressed( npz_path, **arrays, classes=np.array(mapping["classes"], dtype=object), meta_json=np.array(json.dumps(all_meta, default=str)), skipped_json=np.array(json.dumps(all_skipped, default=str)), ) summary = { "protocol": protocol, "sequence_length": config.SEQUENCE_LENGTH, "feature_dim": config.FEATURE_DIM, "model_input_dim": config.MODEL_INPUT_DIM, "mask_enabled": config.MASK_ENABLED, "augment_copies": augment_copies, "seed": seed, "label_sha1": mapping["sha1"], "num_classes": mapping["num_classes"], "shapes": {k: list(v.shape) for k, v in arrays.items()}, "skipped": {k: len(v) for k, v in all_skipped.items()}, "skipped_detail": all_skipped, "seconds": round(time.time() - t0, 2), "normalisation": { "origin": "mid-shoulder (fallback: mid-hip, then landmark centroid)", "scale": "shoulder width (fallback: hip width, then 1.0)", "x_axis_scale": config.X_AXIS_SCALE, "axis_correction_applied": config.APPLY_KEYPOINT_AXIS_CORRECTION, }, "sequence": {"resample": config.SEQ_RESAMPLE_METHOD, "pad_mode": config.SEQ_PAD_MODE, "min_valid_frames": config.MIN_VALID_FRAMES}, "augmentation": dict(config.AUGMENTATION), "split_summary": split_summary.as_dict(), } (out_dir / f"preprocess_summary_{protocol}.json").write_text(json.dumps(summary, indent=2)) (out_dir / "label_mapping.json").write_text(json.dumps(mapping, indent=2)) LOG.info("wrote %s (%.1f MB)", npz_path, npz_path.stat().st_size / 1e6) return npz_path def load_processed(protocol: Optional[str] = None, path: Optional[Path] = None) -> Dict[str, np.ndarray]: protocol = protocol or config.SPLIT_PROTOCOL path = Path(path or (Path(config.PROCESSED_DIR) / f"sequences_{protocol}.npz")) if not path.exists(): raise FileNotFoundError(f"{path} not found - run: python training/preprocess.py --protocol {protocol}") data = np.load(path, allow_pickle=True) return {k: data[k] for k in data.files} def main(argv: Optional[Sequence[str]] = None) -> int: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--protocol", default=config.SPLIT_PROTOCOL, choices=["session_disjoint", "official", "hf"]) ap.add_argument("--augment-copies", type=int, default=_def_copies() if False else None, help="extra augmented copies of each training clip (default from config)") ap.add_argument("--seed", type=int, default=config.RANDOM_SEED) ap.add_argument("--limit", type=int, default=None, help="limit clips per split (smoke test)") ap.add_argument("--out-dir", default=None) args = ap.parse_args(argv) logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") copies = args.augment_copies if copies is None: copies = int(config.AUGMENTATION.get("augment_copies", 1)) if config.AUGMENTATION.get("enabled", True) else 0 try: run(args.protocol, copies, args.seed, args.limit, Path(args.out_dir) if args.out_dir else None) except Exception: traceback.print_exc() return 1 return 0 def _def_copies(): # pragma: no cover - helper for argparse default return None if __name__ == "__main__": raise SystemExit(main())