73 lines
2.7 KiB
Python
73 lines
2.7 KiB
Python
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
COMFY_ROOT = REPO_ROOT.parents[1]
|
|
sys.path.insert(0, str(COMFY_ROOT))
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
|
|
from camera_drive import drive_camera, parse_keyframes, sample_path, _prepare_path # noqa: E402
|
|
|
|
|
|
def person():
|
|
vertices = np.array([
|
|
[-0.5, -1.0, 0.0], [0.5, -1.0, 0.0],
|
|
[-0.5, 1.0, 0.0], [0.5, 1.0, 0.0],
|
|
], dtype=np.float32)
|
|
keypoints = np.zeros((70, 3), dtype=np.float32)
|
|
face = np.zeros((5, 3), dtype=np.float32)
|
|
return {
|
|
"pred_cam_t": np.array([0.0, 0.0, 5.0], dtype=np.float32),
|
|
"pred_vertices": vertices,
|
|
"pred_keypoints_3d": keypoints,
|
|
"pred_keypoints_2d": np.full((70, 2), 50.0, dtype=np.float32),
|
|
"pred_face_keypoints_3d": face,
|
|
"pred_face_keypoints_2d": np.full((5, 2), 50.0, dtype=np.float32),
|
|
"pred_joint_coords": np.zeros((127, 3), dtype=np.float32),
|
|
"focal_length": np.asarray(50.0, dtype=np.float32),
|
|
}
|
|
|
|
|
|
class CameraDriveTests(unittest.TestCase):
|
|
def test_keyframe_path_wraps_across_azimuth_seam(self):
|
|
keyframes = parse_keyframes(
|
|
'[{"f":1,"az":170,"el":0,"dist":1},{"f":3,"az":-170,"el":0,"dist":1}]',
|
|
3,
|
|
)
|
|
azimuth, elevation, distance = sample_path(_prepare_path(keyframes), 2, "linear", False)
|
|
self.assertAlmostEqual(abs(azimuth), 180.0)
|
|
self.assertEqual(elevation, 0.0)
|
|
self.assertEqual(distance, 1.0)
|
|
|
|
def test_static_orbit_reprojects_without_mutating_input(self):
|
|
source_person = person()
|
|
pose = {
|
|
"frames": [[source_person], [person()]],
|
|
"image_size": (100, 100),
|
|
"faces": np.zeros((0, 3), dtype=np.int64),
|
|
}
|
|
out = drive_camera(pose, 45.0, 10.0, 1.0, 0.0, False, 0, "", "linear", "linear")
|
|
result = out["frames"][0][0]
|
|
np.testing.assert_allclose(result["pred_cam_t"], 0.0)
|
|
self.assertFalse(np.allclose(result["pred_vertices"], source_person["pred_vertices"]))
|
|
np.testing.assert_allclose(pose["frames"][0][0]["pred_cam_t"], [0.0, 0.0, 5.0])
|
|
|
|
def test_keyframes_drive_different_viewpoints(self):
|
|
pose = {
|
|
"frames": [[person()], [person()], [person()]],
|
|
"image_size": (100, 100),
|
|
}
|
|
keyframes = '[{"f":1,"az":-45,"el":0,"dist":1},{"f":3,"az":45,"el":0,"dist":1}]'
|
|
out = drive_camera(pose, 0.0, 0.0, 1.0, 0.0, True, 3, keyframes, "linear", "linear")
|
|
first = out["frames"][0][0]["pred_vertices"]
|
|
last = out["frames"][2][0]["pred_vertices"]
|
|
self.assertFalse(np.allclose(first, last))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|