| |
| """ |
| 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_<protocol>.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 |
| from backend.preprocessing import include50, landmarks as lm, sequence as seq_ops |
|
|
| LOG = logging.getLogger("preprocess") |
|
|
|
|
| |
| |
| |
| 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(): |
| return None |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
