127 lines
4.4 KiB
Python
127 lines
4.4 KiB
Python
import json
|
|
import math
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
import torch
|
|
|
|
|
|
folder_paths = types.ModuleType("folder_paths")
|
|
folder_paths.get_input_directory = lambda: "/path/that/does/not/exist"
|
|
folder_paths.get_temp_directory = lambda: "/private/tmp"
|
|
folder_paths.get_annotated_filepath = lambda value: value
|
|
sys.modules.setdefault("folder_paths", folder_paths)
|
|
|
|
from py.nodes.mask_canvas_editor import MaskCanvasEditor # noqa: E402
|
|
|
|
|
|
def grid_to_source(grid, image_height, image_width):
|
|
source_x = (grid[..., 0] + 1.0) * max(image_width - 1, 1) / 2.0
|
|
source_y = (grid[..., 1] + 1.0) * max(image_height - 1, 1) / 2.0
|
|
return source_x, source_y
|
|
|
|
|
|
class SamplingGridTests(unittest.TestCase):
|
|
def test_identity_samples_centered_bbox(self):
|
|
grid = MaskCanvasEditor._build_sampling_grid(
|
|
bbox_h=4,
|
|
bbox_w=6,
|
|
img_h=8,
|
|
img_w=10,
|
|
y_min=3,
|
|
y_max=7,
|
|
x_min=2,
|
|
x_max=8,
|
|
scale=1.0,
|
|
rotation=0.0,
|
|
flip_h=False,
|
|
flip_v=False,
|
|
offset_x=0.0,
|
|
offset_y=0.0,
|
|
device=torch.device("cpu"),
|
|
)
|
|
source_x, source_y = grid_to_source(grid, 8, 10)
|
|
|
|
torch.testing.assert_close(source_x[0, 0], torch.arange(2.0, 8.0))
|
|
torch.testing.assert_close(source_y[0, :, 0], torch.arange(2.0, 6.0))
|
|
|
|
def test_subpixel_offset_is_not_rounded(self):
|
|
grid = MaskCanvasEditor._build_sampling_grid(
|
|
bbox_h=4,
|
|
bbox_w=6,
|
|
img_h=8,
|
|
img_w=10,
|
|
y_min=3,
|
|
y_max=7,
|
|
x_min=2,
|
|
x_max=8,
|
|
scale=2.0,
|
|
rotation=0.0,
|
|
flip_h=False,
|
|
flip_v=False,
|
|
offset_x=1.25,
|
|
offset_y=-0.5,
|
|
device=torch.device("cpu"),
|
|
)
|
|
source_x, source_y = grid_to_source(grid, 8, 10)
|
|
|
|
self.assertAlmostEqual(source_x[0, 2, 3].item(), 4.125, places=5)
|
|
self.assertAlmostEqual(source_y[0, 2, 3].item(), 4.0, places=5)
|
|
|
|
def test_inverse_chain_matches_documented_forward_transform(self):
|
|
scale = 1.7
|
|
rotation = 31.0
|
|
offset_x = 2.4
|
|
offset_y = -1.3
|
|
source_relative_x = -1.2
|
|
source_relative_y = 0.75
|
|
|
|
flipped_x = -source_relative_x
|
|
flipped_y = source_relative_y
|
|
scaled_x = flipped_x * scale
|
|
scaled_y = flipped_y * scale
|
|
radians = math.radians(rotation)
|
|
output_x = scaled_x * math.cos(radians) - scaled_y * math.sin(radians) + offset_x
|
|
output_y = scaled_x * math.sin(radians) + scaled_y * math.cos(radians) + offset_y
|
|
|
|
translated_x = output_x - offset_x
|
|
translated_y = output_y - offset_y
|
|
rotated_x = translated_x * math.cos(radians) + translated_y * math.sin(radians)
|
|
rotated_y = -translated_x * math.sin(radians) + translated_y * math.cos(radians)
|
|
recovered_x = (rotated_x / scale) * -1.0
|
|
recovered_y = rotated_y / scale
|
|
|
|
self.assertAlmostEqual(recovered_x, source_relative_x, places=7)
|
|
self.assertAlmostEqual(recovered_y, source_relative_y, places=7)
|
|
|
|
|
|
class NodeContractTests(unittest.TestCase):
|
|
def test_internal_image_always_has_none_choice(self):
|
|
choices = MaskCanvasEditor.INPUT_TYPES()["optional"]["internal_image"][0]
|
|
self.assertEqual(choices[0], "None")
|
|
|
|
def test_process_accepts_fractional_offsets(self):
|
|
image = torch.linspace(0.0, 1.0, 25).reshape(1, 5, 5, 1).repeat(1, 1, 1, 3)
|
|
mask = torch.ones((1, 3, 3), dtype=torch.float32)
|
|
state = json.dumps({"offsetX": 0.25, "offsetY": -0.5})
|
|
|
|
with mock.patch.object(MaskCanvasEditor, "_save_preview", return_value=("preview.png", 5, 5)):
|
|
result = MaskCanvasEditor().process(mask, image, state, "1")
|
|
|
|
output_image, output_mask = result["result"]
|
|
self.assertEqual(tuple(output_image.shape), (1, 3, 3, 3))
|
|
self.assertEqual(tuple(output_mask.shape), (1, 3, 3))
|
|
self.assertGreater(output_image[0, 1, 1, 0].item(), 0.0)
|
|
|
|
def test_invalid_transform_json_has_clear_error(self):
|
|
image = torch.zeros((1, 2, 2, 3), dtype=torch.float32)
|
|
mask = torch.ones((1, 2, 2), dtype=torch.float32)
|
|
with self.assertRaisesRegex(ValueError, "valid JSON"):
|
|
MaskCanvasEditor().process(mask, image, "not-json", "1")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|