File size: 5,458 Bytes
458c218 | 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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | """MotionDNA v0 — modelo generativo de movimiento humano entrenado con la pose
extraida de music videos. Transformer causal que predice la pose del frame t+1
dada la secuencia 0..t. Genera movimiento nuevo por rollout autoregresivo.
Pensado para caber holgado en una L4 (23GB)."""
import os, glob, math, time, json
import numpy as np
import torch, torch.nn as nn
from torch.utils.data import Dataset, DataLoader
POSES = os.path.expanduser("~/motiondna/poses")
CKPT = os.path.expanduser("~/motiondna/checkpoints")
SAMP = os.path.expanduser("~/motiondna/samples")
os.makedirs(CKPT, exist_ok=True); os.makedirs(SAMP, exist_ok=True)
DEV = "cuda" if torch.cuda.is_available() else "cpu"
WIN = 64 # ventana temporal
DIM = 34 # 17 keypoints * 2
D_MODEL = 256; N_HEAD = 8; N_LAYER = 6; DROP = 0.1
BATCH = 64; EPOCHS = 60; LR = 3e-4
def normalize(seq):
"""seq [T,17,2] xy normalizado a frame -> centrado en cadera + escala torso."""
s = seq.copy().astype(np.float32)
hip = (s[:,11,:] + s[:,12,:]) / 2.0 # centro caderas
s = s - hip[:,None,:]
sh = (s[:,5,:] + s[:,6,:]) / 2.0 # centro hombros
torso = np.linalg.norm(sh - 0.0, axis=1) # dist hombros->cadera(0)
scale = np.clip(np.median(torso), 1e-3, None)
s = s / scale
return s.reshape(s.shape[0], -1) # [T,34]
class MotionDS(Dataset):
def __init__(self):
self.windows = []
files = sorted(glob.glob(os.path.join(POSES, "*.npy")))
for f in files:
seq = np.load(f)
if seq.shape[0] < WIN + 1: continue
n = normalize(seq)
for i in range(0, n.shape[0] - WIN - 1, WIN // 2):
self.windows.append(n[i:i+WIN+1])
if self.windows:
allp = np.concatenate(self.windows, 0)
self.mean = allp.mean(0, keepdims=True); self.std = allp.std(0, keepdims=True) + 1e-6
self.windows = [(w - self.mean) / self.std for w in self.windows]
print(f"Ventanas de entrenamiento: {len(self.windows)} | clips: {len(files)}", flush=True)
def __len__(self): return len(self.windows)
def __getitem__(self, i):
w = torch.from_numpy(self.windows[i].astype(np.float32))
return w[:-1], w[1:] # input, target (shift 1)
class MotionDNA(nn.Module):
def __init__(self):
super().__init__()
self.inp = nn.Linear(DIM, D_MODEL)
self.pos = nn.Parameter(torch.randn(1, WIN, D_MODEL) * 0.02)
layer = nn.TransformerEncoderLayer(D_MODEL, N_HEAD, D_MODEL*4, DROP, batch_first=True, activation="gelu")
self.enc = nn.TransformerEncoder(layer, N_LAYER)
self.out = nn.Linear(D_MODEL, DIM)
def forward(self, x):
T = x.size(1)
h = self.inp(x) + self.pos[:, :T]
mask = torch.triu(torch.ones(T, T, device=x.device), 1).bool()
h = self.enc(h, mask=mask)
return self.out(h)
def main():
ds = MotionDS()
if len(ds) == 0:
print("ERROR: no hay ventanas. Corre extract_poses.py primero."); return
dl = DataLoader(ds, batch_size=BATCH, shuffle=True, num_workers=4, drop_last=True)
model = MotionDNA().to(DEV)
nparam = sum(p.numel() for p in model.parameters())
print(f"Device: {DEV} | Parametros: {nparam/1e6:.2f}M", flush=True)
opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-4)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, EPOCHS)
lossf = nn.SmoothL1Loss()
best = 1e9; hist = []
for ep in range(1, EPOCHS+1):
model.train(); tot = 0; nb = 0
for x, y in dl:
x, y = x.to(DEV), y.to(DEV)
pred = model(x)
loss = lossf(pred, y)
opt.zero_grad(); loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
tot += loss.item(); nb += 1
sched.step()
avg = tot / max(nb,1); hist.append(avg)
mem = torch.cuda.max_memory_allocated()/1e9 if DEV=="cuda" else 0
print(f"Epoch {ep:03d}/{EPOCHS} loss={avg:.5f} lr={sched.get_last_lr()[0]:.2e} vram={mem:.2f}GB", flush=True)
if avg < best:
best = avg
torch.save({"model": model.state_dict(), "mean": ds.mean, "std": ds.std,
"cfg": {"WIN":WIN,"DIM":DIM,"D_MODEL":D_MODEL,"N_HEAD":N_HEAD,"N_LAYER":N_LAYER},
"epoch": ep, "loss": best}, os.path.join(CKPT, "motiondna_v0_best.pt"))
torch.save({"model": model.state_dict(), "mean": ds.mean, "std": ds.std,
"cfg": {"WIN":WIN,"DIM":DIM,"D_MODEL":D_MODEL,"N_HEAD":N_HEAD,"N_LAYER":N_LAYER},
"epoch": EPOCHS, "loss": hist[-1]}, os.path.join(CKPT, "motiondna_v0_final.pt"))
json.dump(hist, open(os.path.join(CKPT, "loss_history.json"), "w"))
# --- rollout de muestra: genera 120 frames de movimiento nuevo ---
model.eval()
with torch.no_grad():
seed = ds[0][0][:16].unsqueeze(0).to(DEV) # 16 frames semilla
gen = seed.clone()
for _ in range(120):
nxt = model(gen[:, -WIN:])[:, -1:]
gen = torch.cat([gen, nxt], 1)
gen = gen[0].cpu().numpy() * ds.std + ds.mean # des-normaliza
np.save(os.path.join(SAMP, "motiondna_sample.npy"), gen.reshape(-1,17,2))
print(f"LISTO. best_loss={best:.5f} ckpt={CKPT}/motiondna_v0_best.pt sample={SAMP}/motiondna_sample.npy", flush=True)
if __name__ == "__main__":
main()
|