Files
AnyImage/tests/test_addon.py
nan 20b653bb30 Refactor image handling and update foreground processing
- Changed the logic in `remove_background_faces` to keep only faces with all foreground vertices.
- Updated background job to save foreground images with a new naming convention (`color.0001.png`).
- Replaced `copy_frames` with `copy_color_images` in various jobs to ensure consistent image naming.
- Introduced a new function `copy_color_images` to handle the copying of images with a specific naming format.
- Updated the BEN2 model to save output images with a new prefix for better clarity.
- Modified tests to reflect changes in image naming and processing logic.
- Added a new shadeless node group for improved material handling in geometry nodes.
- Enhanced validation scripts to check for new image naming conventions and ensure proper functionality.
2026-08-19 01:45:59 +08:00

744 lines
27 KiB
Python

import importlib
import json
import sys
import tempfile
import threading
import unittest
from pathlib import Path
from types import ModuleType, SimpleNamespace
from unittest.mock import patch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
SOURCE_ROOT = PROJECT_ROOT / "src"
class AddonTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.remove_modules()
def append_status_bar(draw):
draw._owner = "snapform"
cls.status_bar_draws.append(draw)
fake_bpy = ModuleType("bpy")
fake_bpy.props = SimpleNamespace(
BoolProperty=lambda **options: options,
EnumProperty=lambda **options: options,
FloatProperty=lambda **options: options,
IntProperty=lambda **options: options,
PointerProperty=lambda **options: options,
StringProperty=lambda **options: options,
)
fake_bpy.types = SimpleNamespace(
AddonPreferences=object,
Menu=object,
Operator=object,
Panel=object,
PropertyGroup=object,
Scene=type("Scene", (), {}),
STATUSBAR_HT_header=SimpleNamespace(
append=append_status_bar,
remove=lambda draw: cls.status_bar_draws.remove(draw),
),
VIEW3D_MT_object_context_menu=SimpleNamespace(
append=lambda draw: cls.context_menu_draws.append(draw),
remove=lambda draw: cls.context_menu_draws.remove(draw),
),
OUTLINER_MT_object=SimpleNamespace(
append=lambda draw: cls.outliner_menu_draws.append(draw),
remove=lambda draw: cls.outliner_menu_draws.remove(draw),
),
)
cls.registered = []
cls.status_bar_draws = []
cls.context_menu_draws = []
cls.outliner_menu_draws = []
fake_bpy.utils = SimpleNamespace(
register_class=cls.registered.append,
unregister_class=cls.registered.remove,
)
fake_bpy.context = SimpleNamespace(
preferences=SimpleNamespace(addons={}),
)
fake_bpy.app = SimpleNamespace(online_access=True)
fake_bpy.path = SimpleNamespace(abspath=lambda path: path)
fake_bpy.data = SimpleNamespace(
objects=SimpleNamespace(get=lambda _name: None),
)
sys.modules["bpy"] = fake_bpy
sys.path.insert(0, str(SOURCE_ROOT))
cls.paths_before_import = tuple(sys.path)
cls.addon = importlib.import_module("snapform")
cls.image_depth = importlib.import_module(
"snapform.blender.depth_plane"
)
cls.image_mesh = importlib.import_module(
"snapform.blender.surface"
)
cls.paths_after_import = tuple(sys.path)
cls.fake_bpy = fake_bpy
@classmethod
def tearDownClass(cls):
sys.path.remove(str(SOURCE_ROOT))
sys.modules.pop("bpy", None)
cls.remove_modules()
@staticmethod
def remove_modules():
for module_name in tuple(sys.modules):
if module_name == "snapform" or module_name.startswith("snapform."):
sys.modules.pop(module_name)
def setUp(self):
preferences = SimpleNamespace(
storage_root="",
device="CUDA",
)
self.fake_bpy.context.preferences.addons = {
"snapform": SimpleNamespace(preferences=preferences)
}
def tearDown(self):
self.addon.runtime.close_active()
self.registered.clear()
self.status_bar_draws.clear()
self.context_menu_draws.clear()
self.outliner_menu_draws.clear()
self.fake_bpy.data = SimpleNamespace(
objects=SimpleNamespace(get=lambda _name: None),
)
if hasattr(self.fake_bpy.types.Scene, "snapform_settings"):
del self.fake_bpy.types.Scene.snapform_settings
def test_snapform_has_no_runtime_boundary_modules(self):
for name in (
"client.py",
"environment.py",
"storage.py",
"models.py",
"status_bar.py",
):
self.assertFalse((SOURCE_ROOT / "snapform" / name).exists(), name)
self.assertFalse(hasattr(self.addon, "client"))
def test_import_does_not_add_addon_directories_to_sys_path(self):
self.assertEqual(self.paths_after_import, self.paths_before_import)
self.assertFalse(
any(
name == "snapform.server" or name.startswith("snapform.server.")
for name in sys.modules
)
)
def test_job_request_is_not_part_of_blendjob(self):
blendjob = importlib.import_module("blendjob")
self.assertFalse(hasattr(blendjob, "JobRequest"))
def test_job_operators_are_flat_and_use_plain_properties(self):
operators = self.addon.operators
for name in (
"SnapFormJobMixin",
"ImageInputJobOperator",
"GeometryImageJobOperator",
"Da3ImageJobOperator",
"Moge2ImageJobOperator",
"BackgroundRequestMixin",
):
self.assertFalse(hasattr(operators, name), name)
for operator in operators.JOB_OPERATOR_CLASSES:
for name in getattr(operator, "__annotations__", {}):
self.assertFalse(name.startswith("snapform_"), (operator, name))
def test_concrete_job_operator_is_the_blender_operator(self):
operator = self.addon.operators.GenerateMoge2Depth
self.assertEqual(operator.bl_idname, "snapform.generate_moge2_depth")
self.assertIn(self.addon.job_runtime.JobOperatorBase, operator.__mro__)
self.assertIn(self.fake_bpy.types.Operator, operator.__mro__)
def test_default_request_sends_every_operator_property(self):
class Example(self.addon.job_runtime.JobOperatorBase):
value: object()
enabled: object()
operator = Example()
operator.value = 42
operator.enabled = True
self.assertEqual(
operator.request(object()),
{"value": 42, "enabled": True},
)
def test_overridden_request_purely_replaces_default_properties(self):
class Example(self.addon.job_runtime.JobOperatorBase):
value: object()
def request(self, context):
return {"replacement": context.value}
operator = Example()
operator.value = 42
self.assertEqual(
operator.request(SimpleNamespace(value=7)),
{"replacement": 7},
)
def test_runtime_submits_job_type_and_plain_dict(self):
calls = []
submitted = threading.Event()
class Controller:
def submit(self, job_type, parameters):
calls.append((job_type, parameters))
submitted.set()
return {"job_id": "job", "directory": "."}
def cancel(self, _job_id):
pass
def mark_job_complete(self, _job_id):
pass
class Example(self.addon.job_runtime.JobOperatorBase):
job_type = "example"
value: object()
def controller(self, _runtime):
return Controller()
timers = []
window_manager = SimpleNamespace(
progress_begin=lambda *_args: None,
progress_update=lambda _value: None,
event_timer_add=lambda *_args, **_kwargs: timers.append(object()) or timers[-1],
event_timer_remove=lambda _timer: None,
modal_handler_add=lambda _operator: None,
)
context = SimpleNamespace(
window_manager=window_manager,
window=object(),
)
operator = Example()
operator.value = 9
self.assertEqual(operator.execute(context), {"RUNNING_MODAL"})
self.assertTrue(submitted.wait(1.0))
self.assertEqual(calls, [("example", {"value": 9})])
def test_runtime_owns_environment_and_storage_paths(self):
runtime = self.addon.runtime
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
with patch.object(runtime, "storage_root_factory", root):
self.assertEqual(runtime.storage_root(), root.resolve())
self.assertEqual(runtime.environment_directory(), root / ".venv")
self.assertEqual(runtime.install_status_path(), root / "install-status.json")
self.assertEqual(runtime.install_log_path(), root / "install.log")
self.assertEqual(runtime.server_log_path(), root / "server.log")
def test_runtime_checks_its_environment_dictionary_hash(self):
runtime = self.addon.runtime
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
python = root / ".venv" / (
"Scripts/python.exe" if sys.platform == "win32" else "bin/python"
)
python.parent.mkdir(parents=True)
python.touch()
(root / "manifest.json").write_text(
json.dumps({"environment_hash": runtime.environment_hash}),
encoding="utf-8",
)
with patch.object(runtime, "storage_root_factory", root):
self.assertTrue(runtime.environment_ready())
def test_runtime_builds_the_standard_installer_command(self):
runtime = self.addon.runtime
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
with patch.object(runtime, "storage_root_factory", root):
command = runtime.install_command()
self.assertEqual(command[0], sys.executable)
self.assertEqual(Path(command[1]).name, "environment.py")
self.assertEqual(Path(command[1]).parent.name, "installer")
self.assertIn("--storage-root", command)
self.assertIn("--config", command)
self.assertNotIn("--venv-dir", command)
self.assertNotIn("--models-dir", command)
self.assertNotIn("--stages", command)
def test_runtime_exposes_declared_environment_and_post_install(self):
runtime = self.addon.runtime
self.assertEqual(runtime.environment["python"], "3.12")
self.assertIn("numpy==2.4.2", runtime.environment["packages"])
self.assertFalse(
any(
"blendjob" in package.lower()
for package in runtime.environment["packages"]
)
)
self.assertTrue(callable(runtime.post_install))
def test_snapform_post_install_is_ordinary_code_using_runtime_request(self):
calls = []
fake_runtime = SimpleNamespace(
request=lambda job_type, parameters: calls.append(
(job_type, parameters)
)
)
self.addon.job_runtime.post_install(fake_runtime)
self.assertEqual(
calls,
[
("download-model", {"model": "MOGE2_VITS_NORMAL"}),
("download-model", {"model": "BEN2_BASE"}),
],
)
def test_snapform_post_install_names_each_failed_model(self):
calls = []
def request(_job_type, parameters):
calls.append(parameters["model"])
raise RuntimeError("network unavailable")
with self.assertRaisesRegex(
RuntimeError,
"MoGe-2 ViT-S Normal: network unavailable; BEN2 Base: network unavailable",
):
self.addon.job_runtime.post_install(
SimpleNamespace(request=request)
)
self.assertEqual(calls, ["MOGE2_VITS_NORMAL", "BEN2_BASE"])
def test_environment_controller_runs_arbitrary_post_install_callable(self):
runtime_module = importlib.import_module("blendjob.runtime")
runtime = self.addon.runtime
calls = []
started = threading.Event()
release = threading.Event()
def post_install(received):
calls.append(received)
started.set()
release.wait(1.0)
controller = runtime_module.EnvironmentController(runtime)
controller.process = SimpleNamespace(poll=lambda: 0, returncode=0)
with tempfile.TemporaryDirectory() as directory, patch.object(
runtime,
"post_install",
post_install,
), patch.object(
runtime,
"install_status_path",
return_value=Path(directory) / "missing.json",
):
self.assertEqual(controller.status("environment")["state"], "running")
self.assertTrue(started.wait(1.0))
release.set()
controller.post_thread.join(1.0)
status = controller.status("environment")
self.assertEqual(calls, [runtime])
self.assertEqual(status["state"], "succeeded")
def test_environment_controller_reports_post_install_failure(self):
runtime_module = importlib.import_module("blendjob.runtime")
runtime = self.addon.runtime
def post_install(_runtime):
raise RuntimeError("BEN2 download failed")
controller = runtime_module.EnvironmentController(runtime)
controller.process = SimpleNamespace(poll=lambda: 0, returncode=0)
with tempfile.TemporaryDirectory() as directory, patch.object(
runtime,
"post_install",
post_install,
), patch.object(
runtime,
"install_status_path",
return_value=Path(directory) / "missing.json",
):
controller.status("environment")
controller.post_thread.join(1.0)
status = controller.status("environment")
self.assertEqual(status["state"], "failed")
self.assertEqual(status["error"], "BEN2 download failed")
def test_post_install_progress_resets_without_stage_fraction(self):
runtime_module = importlib.import_module("blendjob.runtime")
controller = runtime_module.EnvironmentController(self.addon.runtime)
controller._update_post_progress(
{"job_id": "first", "job_type": "download-model", "progress": 1.0}
)
self.assertEqual(controller.post_status["job_id"], "first")
controller._update_post_progress(
{"job_id": "second", "job_type": "download-model", "progress": 0.1}
)
self.assertEqual(controller.post_status["job_id"], "second")
self.assertEqual(controller.post_status["progress"], 0.1)
self.assertNotIn("stage", controller.post_status)
self.assertNotIn("stages", controller.post_status)
def test_conversion_model_lists_default_to_moge2_vits(self):
properties = importlib.import_module("snapform.properties")
catalog = (
{"key": "SMALL", "family": "da3", "enabled": True,
"label": "DA3 Small", "description": "DA3", "license": "Apache"},
{"key": "MOGE2_VITS_NORMAL", "family": "moge2", "enabled": True,
"label": "MoGe-2 ViT-S", "description": "MoGe", "license": "MIT"},
)
with patch.object(properties, "model_catalog", return_value=catalog):
depth_items = properties.depth_plane_model_enum_items(None, None)
surface_items = properties.surface_model_enum_items(None, None)
self.assertEqual(depth_items[0][0], "MOGE2_VITS_NORMAL")
self.assertEqual(surface_items[0][0], "MOGE2_VITS_NORMAL")
def test_required_model_operator_requests_both_models(self):
operator = self.addon.operators.DownloadRequiredModels()
with patch.object(self.addon.operators, "require_environment"):
parameters = operator.request(None)
self.assertEqual(
parameters,
{"models": ["MOGE2_VITS_NORMAL", "BEN2_BASE"]},
)
def test_image_menu_offers_required_model_repair_when_missing(self):
calls = []
class Layout:
operator_context = ""
enabled = True
def __init__(self, enabled=True):
self.enabled = enabled
def row(self):
return Layout(self.enabled)
def operator(self, identifier, **_options):
calls.append((identifier, self.enabled))
return SimpleNamespace()
def separator(self):
calls.append(("separator", self.enabled))
menu = self.addon.SnapFormImageMenu()
menu.layout = Layout()
with patch.object(
self.addon.runtime,
"environment_ready",
return_value=True,
), patch(
"snapform.properties.required_models_ready",
return_value=False,
):
menu.draw(None)
self.assertEqual(
calls[0],
("snapform.download_required_models", True),
)
self.assertIn(("snapform.convert_to_depth_plane", False), calls)
self.assertIn(("snapform.convert_to_surface", False), calls)
self.assertIn(("snapform.remove_image_background", False), calls)
def test_image_menu_keeps_disabled_features_when_environment_is_missing(self):
calls = []
class Layout:
operator_context = ""
enabled = True
def __init__(self, enabled=True):
self.enabled = enabled
def row(self):
return Layout(self.enabled)
def operator(self, identifier, **_options):
calls.append((identifier, self.enabled))
return SimpleNamespace()
def separator(self):
calls.append(("separator", self.enabled))
menu = self.addon.SnapFormImageMenu()
menu.layout = Layout()
with patch.object(
self.addon.runtime,
"environment_ready",
return_value=False,
), patch(
"snapform.properties.required_models_ready",
return_value=False,
):
menu.draw(None)
self.assertEqual(calls[0], ("snapform.install_environment", True))
self.assertEqual(
calls[1],
("snapform.install_environment_mirror", True),
)
self.assertNotIn(("snapform.download_required_models", True), calls)
self.assertIn(("snapform.convert_to_depth_plane", False), calls)
self.assertIn(("snapform.convert_to_surface", False), calls)
self.assertIn(("snapform.remove_image_background", False), calls)
def test_image_operators_require_environment_and_required_models(self):
context = SimpleNamespace(
object=SimpleNamespace(data=SimpleNamespace()),
)
operators = (
(self.image_depth, self.addon.ConvertToDepthPlane),
(self.image_mesh, self.addon.ConvertToSurface),
(self.image_depth, self.addon.RemoveImageBackground),
)
states = (
(False, True, False),
(True, False, False),
(True, True, True),
)
for environment_ready, models_ready, expected in states:
with self.subTest(
environment_ready=environment_ready,
models_ready=models_ready,
expected=expected,
), patch.object(
self.addon.runtime,
"environment_ready",
return_value=environment_ready,
), patch.object(
self.addon.runtime,
"server_busy",
return_value=False,
), patch(
"snapform.properties.required_models_ready",
return_value=models_ready,
):
for module, operator in operators:
with self.subTest(operator=operator.__name__), patch.object(
module,
"is_image_empty",
return_value=True,
), patch.object(
module,
"video_total_frames",
return_value=1,
):
self.assertEqual(operator.poll(context), expected)
def test_ui_reads_required_model_files_without_starting_server(self):
properties = importlib.import_module("snapform.properties")
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
moge = root / "models" / "moge-2-vits-normal-onnx"
ben2 = root / "models" / "BEN2-ONNX" / "onnx"
moge.mkdir(parents=True)
ben2.mkdir(parents=True)
(moge / "model.onnx").touch()
(ben2 / "model_fp16.onnx").touch()
with patch.object(
self.addon.runtime,
"storage_root",
return_value=root,
), patch.object(
self.addon.runtime,
"resource",
side_effect=AssertionError("UI must not start the Server"),
):
self.assertTrue(properties.required_models_ready())
def test_runtime_request_uses_request_response_shape(self):
result = object()
responses = []
with patch.object(
self.addon.runtime.server,
"request",
return_value=result,
) as request:
returned = self.addon.runtime.request(
"example",
{"value": 3},
response=responses.append,
)
self.assertIs(returned, result)
request.assert_called_once_with(
"example",
{"value": 3},
response=responses.append,
progress=None,
cancel_check=None,
)
def test_server_request_returns_job_result_and_calls_response(self):
controller_type = importlib.import_module(
"blendjob.controller"
).ServerController
controller = controller_type.__new__(controller_type)
statuses = iter(
(
{"state": "running", "progress": 0.5},
{
"state": "succeeded",
"progress": 1.0,
"result": {"value": 7},
},
)
)
controller.submit = lambda _job_type, _parameters: {
"job_id": "job",
"directory": ".",
}
controller.status = lambda _job_id: next(statuses)
controller.mark_job_complete = lambda _job_id: None
responses = []
result = controller.request(
"example",
{"input": 3},
response=responses.append,
poll_interval=0.01,
)
self.assertEqual(result.value, {"value": 7})
self.assertEqual(responses, [result])
def test_runtime_keeps_server_as_an_entrypoint_string(self):
path, attribute = self.addon.runtime.server_entrypoint.rsplit(":", 1)
self.assertEqual(Path(path).name, "__init__.py")
self.assertEqual(Path(path).parent.name, "server")
self.assertEqual(attribute, "server")
def test_runtime_registers_system_operators_and_status_bar(self):
self.addon.register()
system_classes = self.addon.runtime.operator_classes()
self.assertTrue(set(system_classes) <= set(self.registered))
self.assertEqual(len(self.status_bar_draws), 1)
self.assertEqual(self.status_bar_draws[0]._owner, "snapform")
self.addon.unregister()
self.assertEqual(self.registered, [])
self.assertEqual(self.status_bar_draws, [])
def test_status_bar_reads_runtime_state(self):
calls = []
class Layout:
def row(self, **_options):
return self
def progress(self, **options):
calls.append(("progress", options))
def operator(self, identifier, **_options):
calls.append(("operator", identifier))
runtime = self.addon.runtime
runtime.active_job = object()
runtime.progress = 0.4
runtime.message = "Working"
runtime._draw_status_bar(SimpleNamespace(layout=Layout()), None)
runtime.active_job = None
self.assertEqual(calls[0][1]["factor"], 0.4)
self.assertEqual(calls[0][1]["text"], "Working")
self.assertEqual(calls[1], ("operator", "snapform.cancel_job"))
def test_moge_depth_request_is_written_on_the_operator(self):
operator = self.addon.operators.GenerateMoge2Depth()
operator.input_path = __file__
operator.source_object_name = ""
operator.model = "MOGE2_VITS_NORMAL"
operator.resolution_level = 6
operator.remove_background = False
settings = SimpleNamespace(video_frames=12)
context = SimpleNamespace(
scene=SimpleNamespace(snapform_settings=settings)
)
with patch.object(
self.addon.operators, "require_environment"
), patch.object(self.addon.operators, "require_model"):
parameters = operator.request(context)
self.assertEqual(parameters["input"], __file__)
self.assertEqual(parameters["model"], "MOGE2_VITS_NORMAL")
self.assertEqual(parameters["resolution_level"], 6)
self.assertNotIn("job_type", parameters)
def test_moge_depth_response_calls_an_ordinary_function(self):
operator = self.addon.operators.GenerateMoge2Depth()
operator.source_object_name = "Reference"
operator.remove_background = False
operator.delete_source = True
operator.shadeless = False
result = object()
context = object()
with patch.object(
self.addon.operators, "build_moge2_depth_response"
) as build:
operator.response(context, result)
build.assert_called_once_with(
context,
result,
source_object_name="Reference",
remove_background=False,
delete_source=True,
shadeless=False,
)
def test_entry_operator_dispatches_plain_properties(self):
operator = self.addon.ConvertToDepthPlane()
operator.model = "MOGE2_VITS_NORMAL"
operator.resolution_level = 7
operator.keep_original = False
operator.shadeless = True
operator.remove_background = False
source_object = SimpleNamespace(
name="Reference",
data=SimpleNamespace(source="FILE"),
)
settings = SimpleNamespace(video_frames=1)
context = SimpleNamespace(
object=source_object,
scene=SimpleNamespace(snapform_settings=settings),
)
calls = []
self.fake_bpy.ops = SimpleNamespace(
snapform=SimpleNamespace(
generate_moge2_depth=lambda **options: calls.append(options)
or {"RUNNING_MODAL"}
)
)
with patch.object(
self.image_depth,
"prepare_image_input",
return_value=(Path(__file__), False),
):
self.assertEqual(operator.execute(context), {"FINISHED"})
self.assertEqual(calls[0]["resolution_level"], 7)
self.assertEqual(calls[0]["source_object_name"], "Reference")
self.assertFalse(any(key.startswith("snapform_") for key in calls[0]))
if __name__ == "__main__":
unittest.main()