120 lines
4.3 KiB
Python
120 lines
4.3 KiB
Python
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
|
|
from shape_fit import ( # noqa: E402
|
|
blend_coefficients,
|
|
build_foreground_masks,
|
|
profile_from_pose_data,
|
|
regenerate_pose_data,
|
|
)
|
|
|
|
|
|
def person(shape_value, scale_value):
|
|
return {
|
|
"shape_params": np.full(45, shape_value, dtype=np.float32),
|
|
"scale_params": np.full(28, scale_value, dtype=np.float32),
|
|
}
|
|
|
|
|
|
class FakeHead:
|
|
def __init__(self):
|
|
self.scale_mean = torch.zeros(68)
|
|
|
|
def mhr_forward(
|
|
self,
|
|
global_trans,
|
|
global_rot,
|
|
body_pose_params,
|
|
hand_pose_params,
|
|
scale_params,
|
|
shape_params,
|
|
expr_params,
|
|
return_keypoints,
|
|
return_joint_coords,
|
|
return_model_params,
|
|
return_joint_rotations,
|
|
):
|
|
batch = shape_params.shape[0]
|
|
verts = torch.zeros(batch, 4, 3)
|
|
verts[:, :, 0] = shape_params[:, :1]
|
|
keypoints = torch.zeros(batch, 308, 3)
|
|
keypoints[:, :, 0] = shape_params[:, :1]
|
|
joints = torch.zeros(batch, 127, 3)
|
|
rotations = torch.eye(3).reshape(1, 1, 3, 3).expand(batch, 127, 3, 3).clone()
|
|
packed_scale = torch.zeros(batch, 68)
|
|
packed_scale[:, :28] = scale_params
|
|
model_params = torch.cat([torch.zeros(batch, 136), packed_scale], dim=1)
|
|
return verts, keypoints, joints, model_params, rotations
|
|
|
|
|
|
class FakeInner:
|
|
def __init__(self):
|
|
self.head_pose = FakeHead()
|
|
|
|
|
|
class ShapeFitTests(unittest.TestCase):
|
|
def test_profile_uses_per_coefficient_median(self):
|
|
pose = {"frames": [[person(1.0, 2.0)], [person(9.0, 8.0)], [person(3.0, 4.0)]]}
|
|
profile = profile_from_pose_data(pose)
|
|
np.testing.assert_allclose(profile["shape_params"], 3.0)
|
|
np.testing.assert_allclose(profile["scale_params"], 4.0)
|
|
|
|
def test_blend_strength_one_replaces_source(self):
|
|
source = np.zeros(3, dtype=np.float32)
|
|
target = np.array([1.0, 2.0, 3.0], dtype=np.float32)
|
|
np.testing.assert_allclose(blend_coefficients(source, target, 1.0), target)
|
|
np.testing.assert_allclose(blend_coefficients(source, target, 0.5), target * 0.5)
|
|
|
|
def test_auto_mask_accepts_foreground_and_comfy_alpha_polarity(self):
|
|
image = torch.zeros(1, 24, 24, 3)
|
|
pose = {"frames": [[{}]]}
|
|
foreground = torch.zeros(1, 24, 24)
|
|
foreground[:, 6:18, 7:17] = 1.0
|
|
comfy_alpha = 1.0 - foreground
|
|
|
|
direct = build_foreground_masks(image, pose, 0, foreground, "auto")[0]
|
|
inverted = build_foreground_masks(image, pose, 0, comfy_alpha, "auto")[0]
|
|
self.assertTrue(direct[12, 12])
|
|
self.assertTrue(inverted[12, 12])
|
|
self.assertFalse(direct[0, 0])
|
|
self.assertFalse(inverted[0, 0])
|
|
|
|
def test_regeneration_updates_cached_fields_without_mutating_input(self):
|
|
base = person(0.0, 0.0)
|
|
base.update({
|
|
"global_rot": np.zeros(3, dtype=np.float32),
|
|
"body_pose_params": np.zeros(133, dtype=np.float32),
|
|
"hand_pose_params": np.zeros(108, dtype=np.float32),
|
|
"expr_params": np.zeros(72, dtype=np.float32),
|
|
"pred_cam_t": np.array([0.0, 0.0, 5.0], dtype=np.float32),
|
|
"focal_length": np.asarray(50.0, dtype=np.float32),
|
|
"pred_keypoints_2d": np.full((70, 2), 12.0, dtype=np.float32),
|
|
"pred_keypoints_3d": np.zeros((70, 3), dtype=np.float32),
|
|
"pred_vertices": np.zeros((4, 3), dtype=np.float32),
|
|
})
|
|
pose = {"frames": [[base]], "image_size": (24, 24), "faces": np.zeros((0, 3), dtype=np.int64)}
|
|
profile = {
|
|
"shape_params": np.full(45, 2.0, dtype=np.float32),
|
|
"scale_params": np.full(28, 3.0, dtype=np.float32),
|
|
}
|
|
|
|
out = regenerate_pose_data(FakeInner(), pose, profile, 1.0, 1.0, 0, "camera", 4)
|
|
result = out["frames"][0][0]
|
|
np.testing.assert_allclose(result["shape_params"], 2.0)
|
|
np.testing.assert_allclose(result["scale_params"], 3.0)
|
|
np.testing.assert_allclose(result["pred_vertices"][:, 0], 2.0)
|
|
np.testing.assert_allclose(result["mhr_model_params"][136:164], 3.0)
|
|
np.testing.assert_allclose(pose["frames"][0][0]["shape_params"], 0.0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|