File size: 2,770 Bytes
723126c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | import json
import re
import numpy as np
from pathlib import Path
from graph.smplx_wholebody import SMPLX_KEEP
ROOT = Path("/home/s56fasla/MOTIONX/data")
OUT = Path("/home/s56fasla/MOTIONX/cache")
SUBSETS = ["animation", "haa500", "humman", "idea400", "kungfu", "music", "perform"]
T = 64
MIN_FRAMES = 16
VAL_FRACTION = 0.05
SEED = 0
CLIP_SUFFIX = re.compile(r"_clip\d+$")
OUT.mkdir(parents=True, exist_ok=True)
rng = np.random.default_rng(SEED)
records = []
for subset in SUBSETS:
jdir = ROOT / "joint" / subset
tdir = ROOT / "semantic_label" / subset
if not jdir.exists():
continue
for npy in sorted(jdir.glob("*.npy")):
txt = tdir / (npy.stem + ".txt")
if not txt.exists():
continue
records.append((subset, npy, txt))
print("paired clips", len(records), flush=True)
groups = {}
for subset, npy, _ in records:
g = subset + "/" + CLIP_SUFFIX.sub("", npy.stem)
groups.setdefault(g, []).append(npy.stem)
keys = sorted(groups)
rng.shuffle(keys)
n_val = int(len(keys) * VAL_FRACTION)
val_groups = set(keys[:n_val])
print("groups", len(keys), "val groups", len(val_groups), flush=True)
kept = []
for subset, npy, txt in records:
g = subset + "/" + CLIP_SUFFIX.sub("", npy.stem)
kept.append((subset, npy, txt, "val" if g in val_groups else "train"))
arr = np.lib.format.open_memmap(OUT / "pose.npy", mode="w+", dtype=np.float16,
shape=(len(kept), T, len(SMPLX_KEEP), 3))
meta = []
idx = np.linspace(0, 1, T)
w = 0
skipped = 0
for i, (subset, npy, txt, split) in enumerate(kept):
try:
a = np.load(npy)
except Exception:
skipped += 1
continue
if a.ndim != 3 or a.shape[0] < MIN_FRAMES or a.shape[1] < 55:
skipped += 1
continue
a = a[:, SMPLX_KEEP, :].astype(np.float32)
pos = idx * (a.shape[0] - 1)
lo = np.floor(pos).astype(np.int64)
hi = np.minimum(lo + 1, a.shape[0] - 1)
frac = (pos - lo)[:, None, None]
out = a[lo] * (1.0 - frac) + a[hi] * frac
out = out - out[0, 0]
arr[w] = out.astype(np.float16)
meta.append({"i": w, "subset": subset, "stem": npy.stem, "split": split,
"frames": int(a.shape[0]), "text": txt.read_text(encoding="utf-8").strip()})
w += 1
if w % 2000 == 0:
print("processed", w, flush=True)
arr.flush()
del arr
full = np.load(OUT / "pose.npy", mmap_mode="r")
np.save(OUT / "pose_trimmed.npy", np.asarray(full[:w]))
(OUT / "pose.npy").unlink()
(OUT / "pose_trimmed.npy").rename(OUT / "pose.npy")
json.dump(meta, open(OUT / "meta.json", "w"))
print("written", w, "skipped", skipped, flush=True)
print("train", sum(1 for m in meta if m["split"] == "train"), "val", sum(1 for m in meta if m["split"] == "val"), flush=True)
|