code/build_text_pose.py

import json
import numpy as np
import torch
import torch.nn.functional as F
from pathlib import Path
from transformers import CLIPTextModelWithProjection, CLIPTokenizer

ROOT = Path("/home/s56fasla/MOTIONX/data")
CACHE = Path("/home/s56fasla/MOTIONX/cache")
CLIP_NAME = "openai/clip-vit-base-patch32"
N_FRAMES = 3
BATCH = 256
FIELDS = ["body_pose", "lhand_pose", "rhand_pose"]

dev = "cuda"
meta = json.load(open(CACHE / "meta.json"))
tok = CLIPTokenizer.from_pretrained(CLIP_NAME)
enc = CLIPTextModelWithProjection.from_pretrained(CLIP_NAME).to(dev).eval()

texts = []
owner = []
mask = np.zeros(len(meta), dtype=bool)
for m in meta:
    p = ROOT / "pose_description" / m["subset"] / (m["stem"] + ".json")
    if not p.exists():
        continue
    try:
        d = json.load(open(p))
    except Exception:
        continue
    keys = sorted(d, key=int)
    if not keys:
        continue
    picks = [keys[0], keys[len(keys) // 2], keys[-1]][:N_FRAMES]
    for k in picks:
        parts = [d[k].get(f, "") for f in FIELDS]
        texts.append(" ".join(x for x in parts if x))
        owner.append(m["i"])
    mask[m["i"]] = True

print("clips with pose_desc", int(mask.sum()), "text rows", len(texts), flush=True)

embs = []
with torch.no_grad():
    for i in range(0, len(texts), BATCH):
        b = tok(texts[i:i + BATCH], padding=True, truncation=True, max_length=77, return_tensors="pt").to(dev)
        e = enc(**b).text_embeds
        embs.append(F.normalize(e, dim=-1).cpu())
        if i % (BATCH * 40) == 0:
            print("encoded", i, "/", len(texts), flush=True)
z = torch.cat(embs).float().numpy()

out = np.zeros((len(meta), z.shape[1]), dtype=np.float32)
counts = np.zeros(len(meta), dtype=np.float32)
np.add.at(out, np.asarray(owner), z)
np.add.at(counts, np.asarray(owner), 1.0)
nz = counts > 0
out[nz] /= counts[nz][:, None]
out[nz] /= np.linalg.norm(out[nz], axis=1, keepdims=True)

np.save(CACHE / "text_emb_pose.npy", out)
np.save(CACHE / "pose_desc_mask.npy", mask)
print("saved", out.shape, "usable", int(mask.sum()), flush=True)