File size: 2,059 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 | 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)
|