diff --git a/scripts/fixtures/take.npz b/scripts/fixtures/take.npz new file mode 100644 index 0000000..d1dd66d Binary files /dev/null and b/scripts/fixtures/take.npz differ diff --git a/scripts/kimodo-resident-worker.py b/scripts/kimodo-resident-worker.py index 27d45f6..307f994 100644 --- a/scripts/kimodo-resident-worker.py +++ b/scripts/kimodo-resident-worker.py @@ -1,7 +1,6 @@ #!/usr/bin/env python3 import argparse import json -import math import os import shutil import sys @@ -15,7 +14,7 @@ import torch from huggingface_hub.constants import HF_HOME from kimodo import DEFAULT_MODEL, load_model -from kimodo.constraints import compute_global_heading, load_constraints_lst +from kimodo.constraints import load_constraints_lst from kimodo.exports.motion_io import save_kimodo_npz from kimodo.meta import parse_prompts_from_meta from kimodo.model.load_model import TEXT_ENCODER_PRESETS @@ -186,20 +185,6 @@ def save_outputs(output, output_base, export_bvh, bvh_standard_tpose): return files -def loop_windows(output): - """One loop window per sample, chosen from Kimodo's own posed joints (always the somaskel77 output skeleton) and the - heading it computes from them. Frame indices into the take, last inclusive.""" - skeleton = model.output_skeleton - hips = skeleton.bone_index["Hips"] - windows = [] - for joints in output["posed_joints"]: - heading = compute_global_heading(torch.from_numpy(joints[None]), skeleton)[0] - angle = torch.atan2(heading[:, 1], heading[:, 0]).numpy() - first, last, pose_cost, velocity_cost = loop_window.loop_window(joints, angle, model.fps, math.ceil(len(joints) / 2), hips) - windows.append({"first": first, "last": last, "poseCost": pose_cost, "velocityCost": velocity_cost}) - return windows - - def save_example_dirs(output_base, texts, num_frames, num_samples, constraints_path): output_stem = os.path.splitext(output_base)[0].rstrip(os.sep) base_name = os.path.basename(output_stem) @@ -335,7 +320,7 @@ def generate(payload): request_bool(payload, "bvh"), request_bool(payload, "bvhStandardTpose") or request_bool(payload, "bvh_standard_tpose"), ) - loops = loop_windows(output) if request_bool(payload, "loop") else None + loops = loop_window.loop_windows(output, model) if request_bool(payload, "loop") else None if request_bool(payload, "saveExampleDir") or request_bool(payload, "save_example_dir"): files.extend(save_example_dirs(output_base, texts, num_frames, num_samples, constraints_path)) generation_count += 1 diff --git a/scripts/loop_window.py b/scripts/loop_window.py index 8c93a56..8d3c542 100644 --- a/scripts/loop_window.py +++ b/scripts/loop_window.py @@ -3,9 +3,13 @@ A seam is the jump from frame b back to frame a. Its cost is the pose difference plus 0.1 s times the velocity difference (Kovar, Gleicher and Pighin 2002, motion-graph transitions), both as RMS over joints, with each frame taken relative to its own Hips and turned to face a fixed way, so where the take wanders and which way it faces -do not count. Numpy only. +do not count. loop_window is numpy only; loop_windows needs torch and kimodo's heading, and no model weights. """ +import math + import numpy as np +import torch +from kimodo.constraints import compute_global_heading VELOCITY_WEIGHT_SECONDS = 0.1 @@ -39,3 +43,18 @@ def loop_window(joints, heading, fps, min_frames, hips=0): best = np.flatnonzero(cost == cost.min()) pick = best[np.lexsort((firsts[best], -(lasts[best] - firsts[best])))[0]] return int(firsts[pick]), int(lasts[pick]), float(pose_cost[pick]), float(velocity_cost[pick]) + + +def loop_windows(output, model): + """One loop window per sample of a Kimodo output, as the worker answers them: {first, last, poseCost, velocityCost}, + frame indices into the take, last inclusive. The joints are posed_joints, always in model.output_skeleton (the + somaskel77, not model.skeleton), and the heading is Kimodo's own from those joints.""" + skeleton = model.output_skeleton + hips = skeleton.bone_index["Hips"] + windows = [] + for joints in output["posed_joints"]: + heading = compute_global_heading(torch.from_numpy(joints[None]), skeleton)[0] + angle = torch.atan2(heading[:, 1], heading[:, 0]).numpy() + first, last, pose_cost, velocity_cost = loop_window(joints, angle, model.fps, math.ceil(len(joints) / 2), hips) + windows.append({"first": first, "last": last, "poseCost": pose_cost, "velocityCost": velocity_cost}) + return windows diff --git a/scripts/test_loop_window.py b/scripts/test_loop_window.py index d12d167..6065729 100644 --- a/scripts/test_loop_window.py +++ b/scripts/test_loop_window.py @@ -1,4 +1,4 @@ -"""python3 -m unittest scripts/test_loop_window.py (numpy only)""" +"""python3 -m unittest scripts/test_loop_window.py (numpy, torch and kimodo; no model weights)""" import math import pathlib import sys @@ -7,7 +7,8 @@ import unittest import numpy as np sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) -from loop_window import loop_window # noqa: E402 +from loop_window import loop_window, loop_windows # noqa: E402 +from types import SimpleNamespace # noqa: E402 FPS = 30 @@ -132,5 +133,57 @@ class LoopWindow(unittest.TestCase): self.assertNotIn("4321", str(caught.exception)) +class StandInSkeleton: + """What kimodo's heading and the worker read of a skeleton: the names and parents of the real one, its Hips index + and the [right, left] hip joints (RightLeg, LeftLeg) that heading is taken from.""" + + def __init__(self, names, parents): + self.bone_order_names = [str(name) for name in names] + self.bone_index = {name: index for index, name in enumerate(self.bone_order_names)} + self.bone_parents = {name: (None if parent < 0 else self.bone_order_names[parent]) for name, parent in zip(self.bone_order_names, parents)} + self.hip_joint_idx = [self.bone_index["RightLeg"], self.bone_index["LeftLeg"]] + + +class WorkerLoopWindowsTest(unittest.TestCase): + """loop_windows on a real generated take (the somaskel77 joints of a live Vili take, 240 frames at 30 fps), with + stand-ins for the model's two skeletons: the 77-joint output skeleton and the 30-joint one the model also carries.""" + + @classmethod + def setUpClass(cls): + data = np.load(pathlib.Path(__file__).resolve().parent / "fixtures" / "take.npz") + cls.joints = data["joints"] + cls.model = SimpleNamespace( + output_skeleton=StandInSkeleton(data["names77"], data["parents77"]), + skeleton=StandInSkeleton(data["names30"], data["parents30"]), + fps=FPS, + ) + + def window(self, joints): + (answer,) = loop_windows({"posed_joints": [joints]}, self.model) + return answer["first"], answer["last"] + + def test_the_real_take_gets_the_window_of_its_own_heading_in_the_output_skeleton(self): + self.assertEqual(self.joints.shape, (240, 77, 3)) + self.assertEqual(self.model.output_skeleton.hip_joint_idx, [72, 67]) + self.assertEqual(self.model.skeleton.hip_joint_idx, [26, 22]) + self.assertEqual(self.window(self.joints), (35, 222)) + # The 30-joint skeleton's hip indices read the wrong joints of a 77-joint take and move the window. + wrong = SimpleNamespace(output_skeleton=self.model.skeleton, fps=FPS) + self.assertNotEqual(loop_windows({"posed_joints": [self.joints]}, wrong)[0]["first"], 35) + + def test_turning_the_body_about_its_hips_by_a_smooth_angle_leaves_the_window(self): + turned = np.stack([ + turn(self.joints[f] - self.joints[f, 0], 0.8 * math.sin(0.05 * f)) + self.joints[f, 0] for f in range(len(self.joints)) + ]).astype(np.float32) + self.assertEqual(self.window(turned), (35, 222)) + + def test_each_sample_gets_its_own_window_with_its_costs(self): + still = np.repeat(self.joints[:1], 240, axis=0) + answers = loop_windows({"posed_joints": [self.joints, still]}, self.model) + self.assertEqual([(a["first"], a["last"]) for a in answers], [(35, 222), (0, 239)]) + self.assertEqual(set(answers[0]), {"first", "last", "poseCost", "velocityCost"}) + self.assertAlmostEqual(answers[1]["poseCost"], 0.0, places=6) + + if __name__ == "__main__": unittest.main()