File size: 1,767 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 | """MotionDNA — extrae secuencias de pose 2D (17 keypoints COCO) de los clips de
music video usando YOLO-pose. Guarda un .npy [T,17,2] (xy normalizado) por clip."""
import os, glob, sys
import numpy as np
import cv2
from ultralytics import YOLO
VIDEOS = os.path.expanduser("~/cosmos-predict2/datasets/boostify_mv_demo/videos")
OUT = os.path.expanduser("~/motiondna/poses")
os.makedirs(OUT, exist_ok=True)
model = YOLO("yolov8n-pose.pt") # se descarga solo la 1a vez
clips = sorted(glob.glob(os.path.join(VIDEOS, "*.mp4")))
print(f"Clips: {len(clips)}", flush=True)
def biggest_person(res):
if res.keypoints is None or res.keypoints.xyn is None: return None
kp = res.keypoints.xyn.cpu().numpy() # [n,17,2]
if kp.shape[0] == 0: return None
if res.boxes is not None and res.boxes.conf is not None and len(res.boxes.conf)==kp.shape[0]:
idx = int(res.boxes.conf.cpu().numpy().argmax())
else:
idx = 0
return kp[idx] # [17,2]
total_frames = 0
for ci, path in enumerate(clips):
cap = cv2.VideoCapture(path)
seq = []; last = None
while True:
ok, frame = cap.read()
if not ok: break
res = model.predict(frame, verbose=False, device=0)[0]
kp = biggest_person(res)
if kp is None:
if last is None: continue
kp = last
last = kp
seq.append(kp.astype(np.float32))
cap.release()
if len(seq) >= 64:
arr = np.stack(seq) # [T,17,2]
np.save(os.path.join(OUT, os.path.basename(path).replace(".mp4",".npy")), arr)
total_frames += len(seq)
print(f" [{ci+1}/{len(clips)}] {os.path.basename(path)} -> {len(seq)} frames", flush=True)
print(f"LISTO. Secuencias guardadas. Frames totales: {total_frames}", flush=True)
|