95 lines
2.7 KiB
Python
95 lines
2.7 KiB
Python
import unittest
|
|
from unittest import mock
|
|
|
|
try:
|
|
from .test_pro_edit_helpers import import_plugin_submodule, load_plugin_package
|
|
except ImportError:
|
|
from test_pro_edit_helpers import import_plugin_submodule, load_plugin_package
|
|
|
|
|
|
class _Expr:
|
|
def __init__(self, text):
|
|
self.text = text
|
|
|
|
def clamp(self, _min_value, _max_value):
|
|
return self
|
|
|
|
def __mul__(self, other):
|
|
other_text = getattr(other, "text", repr(other))
|
|
return _Expr(f"({self.text}*{other_text})")
|
|
|
|
def __rmul__(self, other):
|
|
other_text = getattr(other, "text", repr(other))
|
|
return _Expr(f"({other_text}*{self.text})")
|
|
|
|
def __add__(self, other):
|
|
other_text = getattr(other, "text", repr(other))
|
|
return _Expr(f"({self.text}+{other_text})")
|
|
|
|
def __rsub__(self, other):
|
|
other_text = getattr(other, "text", repr(other))
|
|
return _Expr(f"({other_text}-{self.text})")
|
|
|
|
|
|
class _FakeTorchTensor:
|
|
def __init__(self, label):
|
|
self.label = label
|
|
self.text = label
|
|
|
|
def masked_fill(self, _mask, _fill_value):
|
|
raise AssertionError("torch tensor branch should not call masked_fill directly")
|
|
|
|
def __mul__(self, other):
|
|
other_text = getattr(other, "text", repr(other))
|
|
return _Expr(f"({self.label}*{other_text})")
|
|
|
|
|
|
class _FakeMask:
|
|
def __init__(self, label):
|
|
self.label = label
|
|
|
|
def unsqueeze(self, dim):
|
|
return _Expr(f"{self.label}.unsqueeze({dim})")
|
|
|
|
|
|
class _FakeTorchModule:
|
|
Tensor = _FakeTorchTensor
|
|
|
|
|
|
class _FallbackTensor:
|
|
def __init__(self):
|
|
self.called = False
|
|
|
|
def masked_fill(self, _mask, _fill_value):
|
|
self.called = True
|
|
return "fallback-masked"
|
|
|
|
|
|
class TestProEditUtils(unittest.TestCase):
|
|
def setUp(self):
|
|
plugin = load_plugin_package()
|
|
self.utils = import_plugin_submodule(plugin, "pro_edit.pro_edit_utils")
|
|
|
|
def test_build_masked_pixel_image_prefers_torch_tensor_broadcast_path(self):
|
|
image = _FakeTorchTensor("image")
|
|
mask = _FakeMask("mask")
|
|
|
|
with mock.patch.object(self.utils, "torch", _FakeTorchModule()):
|
|
result = self.utils.build_masked_pixel_image(image, mask, fill_value=0.0)
|
|
|
|
self.assertEqual(result.text, "((image*(1.0-mask.unsqueeze(-1)))+(0.0*mask.unsqueeze(-1)))")
|
|
|
|
def test_build_masked_pixel_image_falls_back_to_masked_fill_for_non_torch_objects(self):
|
|
image = _FallbackTensor()
|
|
mask = object()
|
|
|
|
with mock.patch.object(self.utils, "torch", None):
|
|
result = self.utils.build_masked_pixel_image(image, mask, fill_value=0.5)
|
|
|
|
self.assertEqual(result, "fallback-masked")
|
|
self.assertTrue(image.called)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|