360 lines
14 KiB
Python
360 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from scail2.geometry import frame_bboxes
|
|
from scail2.reference_alignment import (
|
|
SCAIL_POSE2_REFERENCE_GEOMETRY_ALIGNED_ATTR,
|
|
align_reference_image_geometry,
|
|
align_reference_image_geometry_slots,
|
|
)
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
PACKAGE_NAME = "ComfyUI_SCAIL_Pose2_ReferenceAlignmentTestPackage"
|
|
|
|
|
|
def _upper_bbox(tensor, *, rows: int):
|
|
return frame_bboxes(tensor[:, :rows], kind="semantic_rgb_mask")[0]
|
|
|
|
|
|
def import_root_package():
|
|
for name in list(sys.modules):
|
|
if name == PACKAGE_NAME or name.startswith(f"{PACKAGE_NAME}."):
|
|
del sys.modules[name]
|
|
|
|
spec = importlib.util.spec_from_file_location(
|
|
PACKAGE_NAME,
|
|
ROOT / "__init__.py",
|
|
submodule_search_locations=[str(ROOT)],
|
|
)
|
|
assert spec is not None
|
|
assert spec.loader is not None
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class Scail2ReferenceAlignmentTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
try:
|
|
import torch
|
|
except ModuleNotFoundError as exc:
|
|
self.skipTest(f"torch unavailable: {exc}")
|
|
self.torch = torch
|
|
|
|
def _reference_inputs(self):
|
|
torch = self.torch
|
|
ref_image = torch.zeros((1, 10, 6, 3), dtype=torch.float32)
|
|
ref_mask = torch.ones((1, 10, 6, 3), dtype=torch.float32)
|
|
pose_video_mask = torch.zeros((3, 8, 8, 3), dtype=torch.float32)
|
|
|
|
ref_image[0, 1:9, 2:4, 0] = 0.8
|
|
ref_image[0, 1:9, 2:4, 1] = 0.2
|
|
ref_mask[0, 1:9, 2:4, :] = self.torch.tensor([0.0, 0.0, 1.0])
|
|
pose_video_mask[0, 2:8, 3:5, 2] = 1.0
|
|
pose_video_mask[1, 1:7, 4:6, 2] = 1.0
|
|
pose_video_mask[2, 2:8, 3:5, 2] = 1.0
|
|
return ref_image, ref_mask, pose_video_mask
|
|
|
|
def test_alignment_places_reference_mask_on_target_bbox_without_canvas_size_widgets(self) -> None:
|
|
ref_image, ref_mask, pose_video_mask = self._reference_inputs()
|
|
|
|
result = align_reference_image_geometry(
|
|
ref_image=ref_image,
|
|
ref_mask=ref_mask,
|
|
pose_video_mask=pose_video_mask,
|
|
fit_mode="contain",
|
|
anchor="bottom_center",
|
|
target_frame_policy="median_bbox",
|
|
)
|
|
|
|
self.assertEqual((1, 8, 8, 3), tuple(result.ref_image.shape))
|
|
self.assertEqual((1, 8, 8, 3), tuple(result.ref_mask.shape))
|
|
self.assertTrue(
|
|
getattr(result.ref_image, SCAIL_POSE2_REFERENCE_GEOMETRY_ALIGNED_ATTR)
|
|
)
|
|
self.assertEqual(
|
|
(3.0, 2.0, 5.0, 8.0),
|
|
frame_bboxes(result.ref_mask, kind="semantic_rgb_mask")[0].to_tuple(),
|
|
)
|
|
self.assertIn("reference_geometry_alignment", result.summary)
|
|
self.assertIn("target_size=8x8", result.summary)
|
|
|
|
def test_auto_upper_subject_control_corrects_local_portrait_offset(self) -> None:
|
|
torch = self.torch
|
|
ref_image = torch.zeros((1, 16, 10, 3), dtype=torch.float32)
|
|
ref_mask = torch.zeros((1, 16, 10, 3), dtype=torch.float32)
|
|
pose_video_mask = torch.zeros((1, 16, 10, 3), dtype=torch.float32)
|
|
|
|
ref_image[0, 0:6, 3:8, 0] = 1.0
|
|
ref_image[0, 7:16, 0:10, 1] = 1.0
|
|
ref_mask[0, 0:6, 3:8, :] = torch.tensor([0.0, 0.0, 1.0])
|
|
ref_mask[0, 7:16, 0:10, :] = torch.tensor([0.0, 0.0, 1.0])
|
|
|
|
pose_video_mask[0, 0:6, 1:6, 2] = 1.0
|
|
pose_video_mask[0, 7:16, 0:10, 2] = 1.0
|
|
|
|
subject_result = align_reference_image_geometry(
|
|
ref_image=ref_image,
|
|
ref_mask=ref_mask,
|
|
pose_video_mask=pose_video_mask,
|
|
fit_mode="contain",
|
|
anchor="bottom_center",
|
|
control_region="subject",
|
|
)
|
|
auto_result = align_reference_image_geometry(
|
|
ref_image=ref_image,
|
|
ref_mask=ref_mask,
|
|
pose_video_mask=pose_video_mask,
|
|
fit_mode="auto",
|
|
anchor="auto",
|
|
control_region="auto",
|
|
)
|
|
|
|
self.assertEqual(
|
|
(3.0, 0.0, 8.0, 6.0),
|
|
_upper_bbox(subject_result.ref_mask, rows=6).to_tuple(),
|
|
)
|
|
self.assertEqual(
|
|
(1.0, 0.0, 6.0, 6.0),
|
|
_upper_bbox(auto_result.ref_mask, rows=6).to_tuple(),
|
|
)
|
|
self.assertEqual("upper_subject", auto_result.control_region)
|
|
self.assertEqual("center", auto_result.effective_anchor)
|
|
self.assertIn("effective_control_region=upper_subject", auto_result.summary)
|
|
|
|
def test_auto_control_region_preserves_tall_full_body_profile(self) -> None:
|
|
torch = self.torch
|
|
ref_image = torch.zeros((1, 20, 6, 3), dtype=torch.float32)
|
|
ref_mask = torch.zeros((1, 20, 6, 3), dtype=torch.float32)
|
|
pose_video_mask = torch.zeros((1, 20, 6, 3), dtype=torch.float32)
|
|
|
|
ref_image[0, 0:20, 2:4, 0] = 1.0
|
|
ref_mask[0, 0:20, 2:4, :] = torch.tensor([0.0, 0.0, 1.0])
|
|
pose_video_mask[0, 0:20, 2:4, 2] = 1.0
|
|
|
|
result = align_reference_image_geometry(
|
|
ref_image=ref_image,
|
|
ref_mask=ref_mask,
|
|
pose_video_mask=pose_video_mask,
|
|
fit_mode="auto",
|
|
anchor="auto",
|
|
control_region="auto",
|
|
)
|
|
|
|
self.assertEqual("subject", result.control_region)
|
|
self.assertEqual("bottom_center", result.effective_anchor)
|
|
self.assertEqual("contain", result.effective_fit_mode)
|
|
|
|
def test_multi_identity_slot_alignment_does_not_use_aggregate_bbox(self) -> None:
|
|
torch = self.torch
|
|
ref_image = torch.zeros((1, 8, 10, 3), dtype=torch.float32)
|
|
ref_mask = torch.zeros((1, 8, 10, 3), dtype=torch.float32)
|
|
pose_video_mask = torch.zeros((1, 8, 10, 3), dtype=torch.float32)
|
|
|
|
ref_image[0, 1:5, 1:3, 0] = 0.25
|
|
ref_image[0, 1:5, 7:9, 1] = 0.75
|
|
ref_mask[0, 1:5, 1:3, :] = torch.tensor([0.0, 0.0, 1.0])
|
|
ref_mask[0, 1:5, 7:9, :] = torch.tensor([1.0, 0.0, 0.0])
|
|
pose_video_mask[0, 2:6, 2:4, :] = torch.tensor([0.0, 0.0, 1.0])
|
|
pose_video_mask[0, 2:6, 5:7, :] = torch.tensor([1.0, 0.0, 0.0])
|
|
|
|
aggregate = align_reference_image_geometry(
|
|
ref_image=ref_image,
|
|
ref_mask=ref_mask,
|
|
pose_video_mask=pose_video_mask,
|
|
fit_mode="contain",
|
|
anchor="bottom_center",
|
|
control_region="subject",
|
|
)
|
|
slots = align_reference_image_geometry_slots(
|
|
ref_image=ref_image,
|
|
ref_mask=ref_mask,
|
|
pose_video_mask=pose_video_mask,
|
|
fit_mode="contain",
|
|
anchor="bottom_center",
|
|
control_region="subject",
|
|
)
|
|
|
|
self.assertEqual((2.0, 2.0, 7.0, 6.0), aggregate.target_bbox.to_tuple())
|
|
self.assertEqual((0, 1), slots.identity_indices)
|
|
self.assertEqual(1, len(slots.additional_ref_images))
|
|
self.assertEqual(1, len(slots.additional_ref_masks))
|
|
self.assertEqual(
|
|
(2.0, 2.0, 4.0, 6.0),
|
|
frame_bboxes(slots.ref_mask, kind="semantic_rgb_mask")[0].to_tuple(),
|
|
)
|
|
self.assertEqual(
|
|
(5.0, 2.0, 7.0, 6.0),
|
|
frame_bboxes(slots.additional_ref_masks[0], kind="semantic_rgb_mask")[
|
|
0
|
|
].to_tuple(),
|
|
)
|
|
self.assertIn("identity_slots=2", slots.summary)
|
|
|
|
def test_alignment_rejects_empty_reference_mask(self) -> None:
|
|
torch = self.torch
|
|
ref_image = torch.zeros((1, 10, 6, 3), dtype=torch.float32)
|
|
ref_mask = torch.ones((1, 10, 6, 3), dtype=torch.float32)
|
|
pose_video_mask = torch.zeros((1, 8, 8, 3), dtype=torch.float32)
|
|
pose_video_mask[0, 2:8, 3:5, 2] = 1.0
|
|
|
|
with self.assertRaisesRegex(ValueError, "ref_mask foreground area ratio"):
|
|
align_reference_image_geometry(
|
|
ref_image=ref_image,
|
|
ref_mask=ref_mask,
|
|
pose_video_mask=pose_video_mask,
|
|
)
|
|
|
|
def test_root_package_does_not_register_reference_alignment_node(self) -> None:
|
|
package = import_root_package()
|
|
|
|
self.assertNotIn(
|
|
"SCAILPose2ReferenceImageGeometryAlign",
|
|
package.NODE_CLASS_MAPPINGS,
|
|
)
|
|
condition_cls = package.NODE_CLASS_MAPPINGS["SCAILPose2SCAIL2Condition"]
|
|
required = condition_cls.INPUT_TYPES()["required"]
|
|
self.assertIn("reference_fit_mode", required)
|
|
self.assertIn("reference_anchor", required)
|
|
self.assertIn("reference_target_frame_policy", required)
|
|
self.assertIn("reference_control_region", required)
|
|
self.assertIn("reference_bbox_margin", required)
|
|
self.assertIn("reference_max_scale", required)
|
|
self.assertIn("reference_min_mask_area_ratio", required)
|
|
self.assertEqual("auto", required["reference_fit_mode"][1]["default"])
|
|
self.assertEqual("auto", required["reference_anchor"][1]["default"])
|
|
self.assertEqual(
|
|
"median_bbox",
|
|
required["reference_target_frame_policy"][1]["default"],
|
|
)
|
|
self.assertEqual("auto", required["reference_control_region"][1]["default"])
|
|
|
|
def test_condition_replacement_auto_aligns_reference_geometry(self) -> None:
|
|
package = import_root_package()
|
|
condition_cls = package.NODE_CLASS_MAPPINGS["SCAILPose2SCAIL2Condition"]
|
|
ref_image, ref_mask, pose_video_mask = self._reference_inputs()
|
|
|
|
condition, = condition_cls().build(
|
|
pose_video_mask=pose_video_mask[:1],
|
|
ref_image=ref_image,
|
|
ref_mask=ref_mask,
|
|
mode="replacement",
|
|
width=8,
|
|
height=8,
|
|
num_frames=1,
|
|
driving_video=self.torch.zeros((1, 8, 8, 3), dtype=self.torch.float32),
|
|
)
|
|
|
|
self.assertEqual("replacement", condition.mode)
|
|
self.assertEqual((1, 8, 8, 3), tuple(condition.ref_image.shape))
|
|
self.assertEqual((1, 8, 8), tuple(condition.ref_mask_indices.shape))
|
|
self.assertTrue(condition.source_kind.endswith(":reference_geometry_aligned"))
|
|
|
|
def test_condition_replacement_auto_splits_multi_identity_reference_slots(self) -> None:
|
|
package = import_root_package()
|
|
condition_cls = package.NODE_CLASS_MAPPINGS["SCAILPose2SCAIL2Condition"]
|
|
torch = self.torch
|
|
ref_image = torch.zeros((1, 8, 10, 3), dtype=torch.float32)
|
|
ref_mask = torch.zeros((1, 8, 10, 3), dtype=torch.float32)
|
|
pose_video_mask = torch.zeros((1, 8, 10, 3), dtype=torch.float32)
|
|
driving_video = torch.zeros((1, 8, 10, 3), dtype=torch.float32)
|
|
|
|
ref_image[0, 1:5, 1:3, 0] = 0.25
|
|
ref_image[0, 1:5, 7:9, 1] = 0.75
|
|
ref_mask[0, 1:5, 1:3, :] = torch.tensor([0.0, 0.0, 1.0])
|
|
ref_mask[0, 1:5, 7:9, :] = torch.tensor([1.0, 0.0, 0.0])
|
|
pose_video_mask[0, 2:6, 2:4, :] = torch.tensor([0.0, 0.0, 1.0])
|
|
pose_video_mask[0, 2:6, 5:7, :] = torch.tensor([1.0, 0.0, 0.0])
|
|
|
|
condition, = condition_cls().build(
|
|
pose_video_mask=pose_video_mask,
|
|
ref_image=ref_image,
|
|
ref_mask=ref_mask,
|
|
mode="replacement",
|
|
width=10,
|
|
height=8,
|
|
num_frames=1,
|
|
driving_video=driving_video,
|
|
)
|
|
|
|
self.assertEqual(1, condition.identity.reference_identity_count)
|
|
self.assertEqual((1,), condition.identity.additional_reference_identity_counts)
|
|
self.assertEqual(2, condition.identity.reference_slot_count)
|
|
self.assertEqual(1, len(condition.additional_references))
|
|
self.assertEqual("generated_identity_1", condition.additional_references[0].source)
|
|
self.assertEqual(
|
|
(2.0, 2.0, 4.0, 6.0),
|
|
frame_bboxes(condition.ref_image, kind="pose_image")[0].to_tuple(),
|
|
)
|
|
self.assertNotIn(
|
|
"multi_identity_reference_slots_under_provisioned",
|
|
condition.identity.warnings,
|
|
)
|
|
|
|
def test_condition_replacement_preserves_user_additional_references_without_duplicate_auto_split(self) -> None:
|
|
package = import_root_package()
|
|
condition_cls = package.NODE_CLASS_MAPPINGS["SCAILPose2SCAIL2Condition"]
|
|
torch = self.torch
|
|
ref_image = torch.zeros((1, 8, 10, 3), dtype=torch.float32)
|
|
ref_mask = torch.zeros((1, 8, 10, 3), dtype=torch.float32)
|
|
pose_video_mask = torch.zeros((1, 8, 10, 3), dtype=torch.float32)
|
|
driving_video = torch.zeros((1, 8, 10, 3), dtype=torch.float32)
|
|
user_extra = torch.zeros((1, 8, 10, 3), dtype=torch.float32)
|
|
|
|
ref_mask[0, 1:5, 1:3, :] = torch.tensor([0.0, 0.0, 1.0])
|
|
ref_mask[0, 1:5, 7:9, :] = torch.tensor([1.0, 0.0, 0.0])
|
|
pose_video_mask[0, 2:6, 2:4, :] = torch.tensor([0.0, 0.0, 1.0])
|
|
pose_video_mask[0, 2:6, 5:7, :] = torch.tensor([1.0, 0.0, 0.0])
|
|
user_extra[0, 2:6, 5:7, :] = torch.tensor([1.0, 0.0, 0.0])
|
|
|
|
condition, = condition_cls().build(
|
|
pose_video_mask=pose_video_mask,
|
|
ref_image=ref_image,
|
|
ref_mask=ref_mask,
|
|
mode="replacement",
|
|
width=10,
|
|
height=8,
|
|
num_frames=1,
|
|
driving_video=driving_video,
|
|
additional_ref_image=["manual_extra"],
|
|
additional_ref_mask=user_extra,
|
|
)
|
|
|
|
self.assertEqual(1, len(condition.additional_references))
|
|
self.assertEqual("manual_extra", condition.additional_references[0].image)
|
|
self.assertEqual("user", condition.additional_references[0].source)
|
|
|
|
def test_condition_animation_skips_reference_geometry_alignment(self) -> None:
|
|
package = import_root_package()
|
|
condition_cls = package.NODE_CLASS_MAPPINGS["SCAILPose2SCAIL2Condition"]
|
|
torch = self.torch
|
|
ref_image = torch.zeros((1, 8, 8, 3), dtype=torch.float32)
|
|
ref_mask = torch.ones((1, 8, 8, 3), dtype=torch.float32)
|
|
pose_video_mask = torch.zeros((1, 8, 8, 3), dtype=torch.float32)
|
|
ref_mask[0, 2:6, 3:5, :] = torch.tensor([0.0, 0.0, 1.0])
|
|
pose_video_mask[0, 1:7, 2:6, 2] = 1.0
|
|
|
|
condition, = condition_cls().build(
|
|
pose_video_mask=pose_video_mask,
|
|
ref_image=ref_image,
|
|
ref_mask=ref_mask,
|
|
mode="animation",
|
|
width=8,
|
|
height=8,
|
|
num_frames=1,
|
|
pose_video=torch.zeros((1, 8, 8, 3), dtype=torch.float32),
|
|
)
|
|
|
|
self.assertIs(condition.ref_image, ref_image)
|
|
self.assertFalse(condition.source_kind.endswith(":reference_geometry_aligned"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|