320 lines
11 KiB
Python
320 lines
11 KiB
Python
import unittest
|
|
|
|
try:
|
|
from .test_mask_draw_helpers import (
|
|
load_plugin_package,
|
|
torch,
|
|
)
|
|
except ImportError:
|
|
from test_mask_draw_helpers import (
|
|
load_plugin_package,
|
|
torch,
|
|
)
|
|
|
|
|
|
HAS_MASK_CREATE_RUNTIME_DEPS = torch is not None
|
|
MASK_CREATE_RUNTIME_DEPS_MESSAGE = "mask create node tests require torch"
|
|
|
|
|
|
def _require_mask_create_runtime_deps():
|
|
if not HAS_MASK_CREATE_RUNTIME_DEPS:
|
|
raise unittest.SkipTest(MASK_CREATE_RUNTIME_DEPS_MESSAGE)
|
|
|
|
|
|
def make_image(width=8, height=8, color=0.25):
|
|
_require_mask_create_runtime_deps()
|
|
return torch.full((1, height, width, 3), float(color), dtype=torch.float32)
|
|
|
|
|
|
def make_mask(width=8, height=8, value=0.0):
|
|
_require_mask_create_runtime_deps()
|
|
return torch.full((1, height, width), float(value), dtype=torch.float32)
|
|
|
|
|
|
class _FakeVAE:
|
|
def encode(self, image):
|
|
batch, height, width, _channels = tuple(image.shape)
|
|
latent_height = max(1, int(height) // 8)
|
|
latent_width = max(1, int(width) // 8)
|
|
return torch.zeros((batch, 4, latent_height, latent_width), dtype=image.dtype, device=image.device)
|
|
|
|
|
|
@unittest.skipUnless(HAS_MASK_CREATE_RUNTIME_DEPS, MASK_CREATE_RUNTIME_DEPS_MESSAGE)
|
|
class TestMaskCreateNode(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
plugin = load_plugin_package()
|
|
cls.node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleMaskCreate"]()
|
|
cls.repair_prepare_node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleRepairPrepare"]()
|
|
|
|
def _create(self, **overrides):
|
|
params = {
|
|
"image_width": 100,
|
|
"image_height": 80,
|
|
"shape_type": "rectangle",
|
|
"coordinate_mode": "percent",
|
|
"center_x": 0.5,
|
|
"center_y": 0.5,
|
|
"width": 0.3,
|
|
"height": 0.25,
|
|
"radius": 0.15,
|
|
"feather": 0.0,
|
|
"blur": 0.0,
|
|
"invert_mask": False,
|
|
"combine_mode": "replace",
|
|
"input_mask": None,
|
|
}
|
|
params.update(overrides)
|
|
return self.node.create_mask(**params)
|
|
|
|
def test_rectangle_percent_mode_creates_expected_bbox_and_area_info(self):
|
|
mask_image, mask, area_info = self._create()
|
|
|
|
self.assertEqual(tuple(mask.shape), (1, 80, 100))
|
|
self.assertEqual(tuple(mask_image.shape), (1, 80, 100, 3))
|
|
self.assertEqual(area_info["image_size"], [100, 80])
|
|
self.assertEqual(area_info["shape_type"], "rectangle")
|
|
self.assertEqual(area_info["coordinate_mode"], "percent")
|
|
self.assertEqual(area_info["bbox"], [35, 30, 65, 50])
|
|
self.assertEqual(area_info["width"], 30.0)
|
|
self.assertEqual(area_info["height"], 20.0)
|
|
self.assertIsNone(area_info["radius"])
|
|
self.assertEqual(area_info["binary_area_px"], 600)
|
|
self.assertAlmostEqual(area_info["effective_area_px"], 600.0, places=4)
|
|
self.assertAlmostEqual(area_info["area_ratio"], 600.0 / 8000.0, places=4)
|
|
self.assertEqual(float(mask[0, 40, 50].item()), 1.0)
|
|
self.assertEqual(float(mask[0, 5, 5].item()), 0.0)
|
|
self.assertEqual(float(mask_image[0, 40, 50, 0].item()), 1.0)
|
|
self.assertEqual(float(mask_image[0, 5, 5, 0].item()), 0.0)
|
|
|
|
def test_square_pixel_mode_uses_width_as_side(self):
|
|
_mask_image, mask, area_info = self.node.create_mask(
|
|
image_width=64,
|
|
image_height=64,
|
|
shape_type="square",
|
|
coordinate_mode="pixel",
|
|
center_x=20.0,
|
|
center_y=30.0,
|
|
width=10.0,
|
|
height=999.0,
|
|
radius=5.0,
|
|
feather=0.0,
|
|
blur=0.0,
|
|
invert_mask=False,
|
|
combine_mode="replace",
|
|
input_mask=None,
|
|
)
|
|
|
|
self.assertEqual(area_info["bbox"], [15, 25, 25, 35])
|
|
self.assertEqual(area_info["binary_area_px"], 100)
|
|
self.assertEqual(area_info["height"], 10.0)
|
|
self.assertEqual(float(mask[0, 30, 20].item()), 1.0)
|
|
self.assertEqual(float(mask[0, 5, 5].item()), 0.0)
|
|
|
|
def test_circle_generates_centered_mask(self):
|
|
_mask_image, mask, area_info = self.node.create_mask(
|
|
image_width=101,
|
|
image_height=101,
|
|
shape_type="circle",
|
|
coordinate_mode="pixel",
|
|
center_x=50.0,
|
|
center_y=50.0,
|
|
width=20.0,
|
|
height=20.0,
|
|
radius=10.0,
|
|
feather=0.0,
|
|
blur=0.0,
|
|
invert_mask=False,
|
|
combine_mode="replace",
|
|
input_mask=None,
|
|
)
|
|
|
|
self.assertEqual(area_info["radius"], 10.0)
|
|
self.assertGreater(area_info["binary_area_px"], 250)
|
|
self.assertLess(area_info["binary_area_px"], 380)
|
|
self.assertEqual(float(mask[0, 50, 50].item()), 1.0)
|
|
self.assertEqual(float(mask[0, 10, 10].item()), 0.0)
|
|
|
|
def test_ellipse_generates_nonempty_mask(self):
|
|
_mask_image, mask, area_info = self.node.create_mask(
|
|
image_width=120,
|
|
image_height=90,
|
|
shape_type="ellipse",
|
|
coordinate_mode="pixel",
|
|
center_x=60.0,
|
|
center_y=45.0,
|
|
width=40.0,
|
|
height=20.0,
|
|
radius=10.0,
|
|
feather=0.0,
|
|
blur=0.0,
|
|
invert_mask=False,
|
|
combine_mode="replace",
|
|
input_mask=None,
|
|
)
|
|
|
|
self.assertEqual(area_info["width"], 40.0)
|
|
self.assertEqual(area_info["height"], 20.0)
|
|
self.assertGreater(area_info["binary_area_px"], 500)
|
|
self.assertEqual(float(mask[0, 45, 60].item()), 1.0)
|
|
self.assertEqual(float(mask[0, 5, 5].item()), 0.0)
|
|
|
|
def test_union_subtract_and_intersect_with_input_mask(self):
|
|
input_mask = make_mask(width=32, height=32, value=1.0)
|
|
|
|
_mask_image, union_mask, union_info = self.node.create_mask(
|
|
image_width=32,
|
|
image_height=32,
|
|
shape_type="rectangle",
|
|
coordinate_mode="pixel",
|
|
center_x=16.0,
|
|
center_y=16.0,
|
|
width=10.0,
|
|
height=10.0,
|
|
radius=5.0,
|
|
feather=0.0,
|
|
blur=0.0,
|
|
invert_mask=False,
|
|
combine_mode="union",
|
|
input_mask=input_mask,
|
|
)
|
|
self.assertEqual(union_info["binary_area_px"], 32 * 32)
|
|
self.assertTrue(torch.all(union_mask == 1.0))
|
|
|
|
_mask_image, subtract_mask, subtract_info = self.node.create_mask(
|
|
image_width=32,
|
|
image_height=32,
|
|
shape_type="rectangle",
|
|
coordinate_mode="pixel",
|
|
center_x=16.0,
|
|
center_y=16.0,
|
|
width=10.0,
|
|
height=10.0,
|
|
radius=5.0,
|
|
feather=0.0,
|
|
blur=0.0,
|
|
invert_mask=False,
|
|
combine_mode="subtract",
|
|
input_mask=input_mask,
|
|
)
|
|
self.assertLess(subtract_info["binary_area_px"], 32 * 32)
|
|
self.assertEqual(float(subtract_mask[0, 16, 16].item()), 0.0)
|
|
|
|
_mask_image, intersect_mask, intersect_info = self.node.create_mask(
|
|
image_width=32,
|
|
image_height=32,
|
|
shape_type="rectangle",
|
|
coordinate_mode="pixel",
|
|
center_x=16.0,
|
|
center_y=16.0,
|
|
width=10.0,
|
|
height=10.0,
|
|
radius=5.0,
|
|
feather=0.0,
|
|
blur=0.0,
|
|
invert_mask=False,
|
|
combine_mode="intersect",
|
|
input_mask=input_mask,
|
|
)
|
|
self.assertEqual(intersect_info["binary_area_px"], 100)
|
|
self.assertEqual(float(intersect_mask[0, 16, 16].item()), 1.0)
|
|
self.assertEqual(float(intersect_mask[0, 1, 1].item()), 0.0)
|
|
|
|
def test_invert_mask_flips_final_mask(self):
|
|
_mask_image, mask, area_info = self._create(invert_mask=True)
|
|
|
|
self.assertEqual(area_info["invert_mask"], True)
|
|
self.assertEqual(float(mask[0, 40, 50].item()), 0.0)
|
|
self.assertEqual(float(mask[0, 5, 5].item()), 1.0)
|
|
|
|
def test_feather_and_blur_create_soft_mask(self):
|
|
_mask_image, mask, area_info = self._create(
|
|
shape_type="rectangle",
|
|
coordinate_mode="pixel",
|
|
center_x=50.0,
|
|
center_y=40.0,
|
|
width=20.0,
|
|
height=20.0,
|
|
feather=4.0,
|
|
blur=3.0,
|
|
)
|
|
|
|
self.assertGreater(area_info["effective_area_px"], 0.0)
|
|
self.assertLess(area_info["effective_area_px"], 8000.0)
|
|
self.assertTrue(bool(((mask > 0.0) & (mask < 1.0)).any().item()))
|
|
|
|
def test_shape_outside_image_returns_black_mask_without_crashing(self):
|
|
_mask_image, mask, area_info = self._create(
|
|
shape_type="circle",
|
|
coordinate_mode="percent",
|
|
center_x=2.0,
|
|
center_y=2.0,
|
|
radius=0.1,
|
|
)
|
|
|
|
self.assertEqual(area_info["binary_area_px"], 0)
|
|
self.assertEqual(area_info["area_ratio"], 0.0)
|
|
self.assertTrue(torch.all(mask == 0.0))
|
|
|
|
def test_input_mask_is_resized_and_output_connects_to_repair_prepare(self):
|
|
image = make_image(width=64, height=64, color=0.25)
|
|
input_mask = make_mask(width=16, height=16, value=1.0)
|
|
|
|
_mask_image, mask_out, area_info = self.node.create_mask(
|
|
image_width=64,
|
|
image_height=64,
|
|
shape_type="rectangle",
|
|
coordinate_mode="pixel",
|
|
center_x=32.0,
|
|
center_y=32.0,
|
|
width=10.0,
|
|
height=10.0,
|
|
radius=5.0,
|
|
feather=0.0,
|
|
blur=0.0,
|
|
invert_mask=False,
|
|
combine_mode="union",
|
|
input_mask=input_mask,
|
|
)
|
|
|
|
latent, work_image, work_mask, repair_info, recommended_denoise, positive, negative = self.repair_prepare_node.prepare(
|
|
image=image,
|
|
mask=mask_out,
|
|
vae=_FakeVAE(),
|
|
repair_scope="region",
|
|
repair_kernel="vae_inpaint",
|
|
task_hint="repair",
|
|
mask_grow=0,
|
|
mask_blur=0.0,
|
|
mask_threshold=0.5,
|
|
invert_mask=False,
|
|
crop_context=64,
|
|
crop_context_factor=1.5,
|
|
min_size=256,
|
|
max_size=1024,
|
|
resize_mode="fit",
|
|
expand_left=0,
|
|
expand_right=0,
|
|
expand_top=0,
|
|
expand_bottom=0,
|
|
canvas_fill="edge",
|
|
auto_recommend="enabled",
|
|
model_info=None,
|
|
positive=None,
|
|
negative=None,
|
|
)
|
|
|
|
self.assertEqual(tuple(mask_out.shape), (1, 64, 64))
|
|
self.assertEqual(tuple(work_image.shape), (1, 64, 64, 3))
|
|
self.assertEqual(tuple(work_mask.shape), (1, 64, 64))
|
|
self.assertEqual(latent["source"], "repair_prepare_region")
|
|
self.assertTrue(repair_info["has_mask"])
|
|
self.assertGreater(area_info["binary_area_px"], 0)
|
|
self.assertEqual(recommended_denoise, 0.45)
|
|
self.assertIsNone(positive)
|
|
self.assertIsNone(negative)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|