Buckets:
| """Cheap ego-motion token: Farneback optical-flow statistics per exterior view per protocol step (CPU). | |
| python aide/features/extract_egomotion.py [--workers 3] [--smoke 5] | |
| For every clip, view in (front, left, right) and protocol step j = 1..9 the dense flow between the | |
| protocol frames j-1 -> j (cache frame_idx10, JPEGs downscaled to 320x180, grayscale) is summarised by | |
| 16 numbers: mean dx, mean dy, mean |f|, std |f|, median |f|, fraction |f| > 1 px, divergence and curl | |
| proxies (mean spatial gradients), and the 2x2-grid means of dx and dy (8). Step 0 copies step 1. | |
| Output results/feats/feats_egomotion.npz: ego [N,10,3,16] f32 in cache clip order, clip_id, views, | |
| frame_idx, feature names; resumable through feats_egomotion_partial.npz. Only frames <= the observed | |
| window are ever handed to a model (the trainer slices [:T_OBS]); this file just stores all 10 steps. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| for _v in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_NUM_THREADS"): | |
| os.environ.setdefault(_v, "1") | |
| import argparse # noqa: E402 | |
| import sys # noqa: E402 | |
| import time # noqa: E402 | |
| from multiprocessing import Pool # noqa: E402 | |
| from pathlib import Path # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| ROOT = Path(__file__).resolve().parents[2] | |
| sys.path.insert(0, str(ROOT)) | |
| from aide.data import protocol as P # noqa: E402 | |
| from common.runmeta import write_run_json # noqa: E402 | |
| EXT_VIEWS = ("front", "left", "right") | |
| SIZE = (320, 180) | |
| NAMES = ["mean_dx", "mean_dy", "mean_mag", "std_mag", "median_mag", "frac_mag_gt1", "div", "curl", | |
| "g00_dx", "g01_dx", "g10_dx", "g11_dx", "g00_dy", "g01_dy", "g10_dy", "g11_dy"] | |
| OUT = P.RESULTS / "feats" / "feats_egomotion.npz" | |
| PARTIAL = P.RESULTS / "feats" / "feats_egomotion_partial.npz" | |
| def flow_stats(prev, cur): | |
| import cv2 | |
| f = cv2.calcOpticalFlowFarneback(prev, cur, None, 0.5, 3, 15, 3, 5, 1.2, 0) | |
| dx, dy = f[..., 0], f[..., 1] | |
| mag = np.sqrt(dx ** 2 + dy ** 2) | |
| div = np.gradient(dx, axis=1).mean() + np.gradient(dy, axis=0).mean() | |
| curl = np.gradient(dy, axis=1).mean() - np.gradient(dx, axis=0).mean() | |
| h, w = dx.shape | |
| grid = [a[i * h // 2:(i + 1) * h // 2, j * w // 2:(j + 1) * w // 2].mean() for a in (dx, dy) for i in (0, 1) for j in (0, 1)] | |
| return np.array([dx.mean(), dy.mean(), mag.mean(), mag.std(), np.median(mag), (mag > 1.0).mean(), div, curl, *grid], np.float32) | |
| def clip_feats(args): | |
| cid, frame_idx = args | |
| import cv2 | |
| cv2.setNumThreads(1) | |
| out = np.zeros((P.N_SAMPLE, len(EXT_VIEWS), len(NAMES)), np.float32) | |
| for vi, v in enumerate(EXT_VIEWS): | |
| prev = None | |
| for j, f in enumerate(frame_idx): | |
| im = cv2.imread(str(P.AIDE_ROOT / cid / f"{v}frames/{int(f)}.jpg"), cv2.IMREAD_GRAYSCALE) | |
| assert im is not None, (cid, v, f) | |
| g = cv2.resize(im, SIZE, interpolation=cv2.INTER_AREA) | |
| if prev is not None: | |
| out[j, vi] = flow_stats(prev, g) | |
| prev = g | |
| out[0, vi] = out[1, vi] | |
| return cid, out | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--workers", type=int, default=3) | |
| ap.add_argument("--smoke", type=int, default=0) | |
| ap.add_argument("--chunk", type=int, default=150) | |
| a = ap.parse_args() | |
| t0 = time.time() | |
| c = np.load(P.CACHE_NPZ, allow_pickle=False) | |
| ids, fidx = c["clip_id"], c["frame_idx10"] | |
| N = len(ids) | |
| ego = np.full((N, P.N_SAMPLE, len(EXT_VIEWS), len(NAMES)), np.nan, np.float32) | |
| done = np.zeros(N, bool) | |
| if PARTIAL.exists() and not a.smoke: | |
| z = np.load(PARTIAL, allow_pickle=False) | |
| assert np.array_equal(z["clip_id"].astype(str), ids.astype(str)) | |
| ego, done = z["ego"], z["done"] | |
| print(f"resume: {done.sum()} clips done") | |
| todo = [i for i in range(N) if not done[i]] | |
| if a.smoke: | |
| todo = todo[:a.smoke] | |
| print(f"{len(todo)} clips to process with {a.workers} workers", flush=True) | |
| pos = {cid: i for i, cid in enumerate(ids)} | |
| with Pool(a.workers) as pool: | |
| for s in range(0, len(todo), a.chunk): | |
| batch = [(ids[i], fidx[i]) for i in todo[s:s + a.chunk]] | |
| for cid, feats in pool.imap_unordered(clip_feats, batch, chunksize=4): | |
| ego[pos[cid]] = feats | |
| done[pos[cid]] = True | |
| el = time.time() - t0 | |
| print(f"[{min(s + a.chunk, len(todo))}/{len(todo)}] {el / 60:.1f} min, {el / max(1, s + len(batch)):.2f} s/clip", flush=True) | |
| if not a.smoke: | |
| np.savez(PARTIAL, ego=ego, done=done, clip_id=ids.astype(str)) | |
| if a.smoke: | |
| sm = ego[[pos[cid] for cid, _ in [(ids[i], None) for i in todo]]] | |
| print("smoke stats (mean over clips/steps/views):", dict(zip(NAMES, np.round(sm.mean((0, 1, 2)), 3).tolist()))) | |
| return | |
| assert done.all() and np.isfinite(ego).all() | |
| np.savez(OUT, ego=ego, clip_id=ids.astype(str), split=c["split"].astype(str), views=np.array(EXT_VIEWS), | |
| frame_idx=fidx.astype(np.int32), names=np.array(NAMES), size=np.array(SIZE)) | |
| stats = {"N": int(N), "shape": list(ego.shape), "bytes": int(OUT.stat().st_size), "size": SIZE, | |
| "mean": dict(zip(NAMES, ego.mean((0, 1, 2)).round(4).tolist())), | |
| "std": dict(zip(NAMES, ego.std((0, 1, 2)).round(4).tolist())), "wall_min": round((time.time() - t0) / 60, 1)} | |
| print(stats) | |
| write_run_json(OUT.parent, "extract_egomotion", [__file__, P.__file__], [P.CACHE_NPZ], seed=None, t_start=t0, | |
| input_listing=False, extra=stats) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 5.6 kB
- Xet hash:
- da3b0e5b7d1b459e5fef76ba653857177700cfdc33da45221c403c29ea4a6e11
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.