nan
e8129fe495
- Added validation documentation for AI model tiers. - Introduced `background.py` for background removal logic, including inference functions for alpha and foreground extraction. - Created `onnx_birefnet.py` to handle BiRefNet model sessions and alpha inference. - Developed tests for background models, ensuring proper dispatch and session reuse. - Added tests for BiRefNet ONNX model to validate alpha output and error handling for invalid values. - Implemented model export scripts for BiRefNet and HAT Sharper, ensuring proper weight compression and licensing. - Created `deform_sampling.py` for deformable convolution operations used in BiRefNet. - Added utility functions for compressing ONNX weights without altering graph arithmetic. - Introduced a reference plane node for depth field adjustments in geometry processing.
414 lines
15 KiB
Python
414 lines
15 KiB
Python
from contextlib import nullcontext
|
|
import importlib
|
|
import json
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
from tests.support.blender import BlenderTestCase
|
|
|
|
|
|
class CutoutOperatorsTest(BlenderTestCase):
|
|
def test_flat_and_solid_without_ai_build_locally(self):
|
|
import numpy as np
|
|
|
|
source = SimpleNamespace(
|
|
data=SimpleNamespace(size=(10, 10)), select_set=lambda _value: None
|
|
)
|
|
context = SimpleNamespace(
|
|
view_layer=SimpleNamespace(objects=SimpleNamespace(active=None))
|
|
)
|
|
selection_mask = self.image_selection.SelectionMask(
|
|
np.ones((8, 8), dtype=np.float32), (1, 1, 9, 9)
|
|
)
|
|
source_rgba = np.ones((10, 10, 4), dtype=np.float32)
|
|
for shape in ("FLAT", "SOLID"):
|
|
with self.subTest(shape=shape):
|
|
operator = SimpleNamespace(
|
|
source_object_name="Source",
|
|
selection_path_json=json.dumps(
|
|
{"points": [[1, 1], [9, 1], [9, 9], [1, 9]]}
|
|
),
|
|
shape=shape,
|
|
point_spacing=32,
|
|
alpha_threshold=0.25,
|
|
fine_outline=False,
|
|
generate_normal=False,
|
|
report=lambda *_args: None,
|
|
_needs_generation=lambda: False,
|
|
)
|
|
with (
|
|
patch.object(
|
|
self.cutout_main,
|
|
"require_image_empty",
|
|
return_value=source,
|
|
),
|
|
patch.object(
|
|
self.cutout_main, "video_total_frames", return_value=1
|
|
),
|
|
patch.object(
|
|
self.cutout_main, "image_rgba", return_value=source_rgba
|
|
) as read_rgba,
|
|
patch.object(
|
|
self.cutout_main,
|
|
"rasterize_selection_path",
|
|
return_value=selection_mask,
|
|
) as rasterize,
|
|
patch.object(
|
|
self.cutout_main,
|
|
"configured_max_ai_input_size",
|
|
return_value=2048,
|
|
),
|
|
patch.object(
|
|
self.cutout_main,
|
|
"configured_max_cutout_mesh_resolution",
|
|
return_value=1024,
|
|
),
|
|
patch.object(
|
|
self.cutout_main,
|
|
"_cutout_content_values",
|
|
return_value=selection_mask.values,
|
|
),
|
|
patch.object(
|
|
self.cutout_main,
|
|
"material_color_image",
|
|
return_value=nullcontext("Color"),
|
|
),
|
|
patch.object(
|
|
self.cutout_main, "create_cutout_shape"
|
|
) as convert,
|
|
):
|
|
result = self.cutout_main.CutoutSelectionToShape.execute(
|
|
operator, context
|
|
)
|
|
|
|
self.assertEqual(result, {"FINISHED"})
|
|
read_rgba.assert_called_once_with(source.data)
|
|
rasterize.assert_called_once()
|
|
convert.assert_called_once_with(
|
|
context,
|
|
source,
|
|
shape,
|
|
32.0,
|
|
selection_mask.values,
|
|
selection_mask.bounds,
|
|
inward_axis="NEGATIVE_Z",
|
|
alpha_threshold=0.25,
|
|
fine_outline=False,
|
|
color_image="Color",
|
|
)
|
|
|
|
|
|
def test_ai_cutout_writes_the_selection_bounds_image(self):
|
|
import numpy as np
|
|
|
|
source = SimpleNamespace(
|
|
data=SimpleNamespace(size=(4096, 2048)),
|
|
select_set=lambda _value: None,
|
|
)
|
|
context = SimpleNamespace(
|
|
view_layer=SimpleNamespace(objects=SimpleNamespace(active=None))
|
|
)
|
|
source_rgba = np.ones((4, 8, 4), dtype=np.float32)
|
|
selection_mask = self.image_selection.SelectionMask(
|
|
np.ones((2, 6), dtype=np.float32),
|
|
(1, 1, 7, 3),
|
|
)
|
|
operator = self.cutout_main.CutoutSelectionToShape()
|
|
operator.alpha_threshold = 0.25
|
|
operator.fine_outline = False
|
|
operator.source_object_name = "Source"
|
|
operator.selection_path_json = json.dumps(
|
|
{"points": [[1, 1], [7, 1], [7, 3], [1, 3]]}
|
|
)
|
|
operator.shape = "DEPTH_SOLID"
|
|
operator.point_spacing = 32
|
|
operator.generate_normal = False
|
|
operator.report = lambda *_args: None
|
|
|
|
with (
|
|
patch.object(
|
|
self.cutout_main,
|
|
"require_image_empty",
|
|
return_value=source,
|
|
),
|
|
patch.object(self.cutout_main, "video_total_frames", return_value=1),
|
|
patch.object(
|
|
self.cutout_main,
|
|
"image_rgba",
|
|
return_value=source_rgba,
|
|
) as read_rgba,
|
|
patch.object(
|
|
self.cutout_main,
|
|
"rasterize_selection_path",
|
|
return_value=selection_mask,
|
|
) as rasterize,
|
|
patch.object(
|
|
self.cutout_main,
|
|
"configured_max_ai_input_size",
|
|
return_value=2048,
|
|
),
|
|
patch.object(
|
|
self.cutout_main,
|
|
"configured_max_cutout_mesh_resolution",
|
|
return_value=1024,
|
|
),
|
|
patch.object(
|
|
self.cutout_main,
|
|
"invoke_ai_setup_if_needed",
|
|
return_value=None,
|
|
),
|
|
patch.object(
|
|
self.cutout_main,
|
|
"cutout_geometry_settings",
|
|
return_value=("MODEL", 5),
|
|
),
|
|
patch.object(
|
|
self.cutout_main,
|
|
"prepare_material_color_input",
|
|
return_value=Path(__file__),
|
|
) as prepare_input,
|
|
patch.object(self.cutout_main, "material_analysis_input", return_value=Path(__file__)),
|
|
patch.object(
|
|
self.cutout_main.JobOperatorBase,
|
|
"execute",
|
|
return_value={"FINISHED"},
|
|
),
|
|
):
|
|
result = operator.execute(context)
|
|
|
|
self.assertEqual(result, {"FINISHED"})
|
|
read_rgba.assert_called_once_with(source.data)
|
|
rasterize.assert_called_once()
|
|
prepare_input.assert_called_once_with(
|
|
source.data,
|
|
selection_mask.bounds,
|
|
source_rgba,
|
|
)
|
|
|
|
|
|
def test_empty_cutout_selection_is_rejected_after_shape_choice(self):
|
|
source = SimpleNamespace(
|
|
data=SimpleNamespace(size=(4096, 2048)),
|
|
select_set=lambda _value: None,
|
|
)
|
|
context = SimpleNamespace(
|
|
view_layer=SimpleNamespace(objects=SimpleNamespace(active=None))
|
|
)
|
|
reports = []
|
|
operator = SimpleNamespace(
|
|
source_object_name="Source",
|
|
selection_path_json=json.dumps(
|
|
{"points": [[1, 1], [7, 1], [7, 3], [1, 3]]}
|
|
),
|
|
shape="FLAT",
|
|
point_spacing=32,
|
|
generate_normal=False,
|
|
report=lambda level, message: reports.append((level, message)),
|
|
)
|
|
|
|
with (
|
|
patch.object(
|
|
self.cutout_main,
|
|
"require_image_empty",
|
|
return_value=source,
|
|
),
|
|
patch.object(self.cutout_main, "video_total_frames", return_value=1),
|
|
patch.object(self.cutout_main, "image_rgba"),
|
|
patch.object(
|
|
self.cutout_main,
|
|
"rasterize_selection_path",
|
|
side_effect=self.image_selection.ImageEditWarning(
|
|
"The selection contains no visible image pixels"
|
|
),
|
|
),
|
|
patch.object(self.cutout_main, "create_cutout_shape") as create_shape,
|
|
):
|
|
result = self.cutout_main.CutoutSelectionToShape.execute(
|
|
operator,
|
|
context,
|
|
)
|
|
|
|
self.assertEqual(result, {"CANCELLED"})
|
|
create_shape.assert_not_called()
|
|
self.assertEqual(
|
|
reports,
|
|
[
|
|
(
|
|
{"WARNING"},
|
|
"The selection contains no visible image pixels",
|
|
)
|
|
],
|
|
)
|
|
|
|
|
|
def test_cutout_request_describes_ai_capabilities(self):
|
|
operator = self.cutout_main.CutoutSelectionToShape()
|
|
operator.alpha_threshold = 0.25
|
|
operator.fine_outline = False
|
|
operator.input_path = __file__
|
|
operator.color_path = __file__
|
|
operator.source_object_name = "Source"
|
|
operator.shape = "DEPTH_SOLID"
|
|
operator.model = self.anyimage.preferences.DEFAULT_GEOMETRY_MODEL_KEY
|
|
operator.resolution_level = 5
|
|
operator.max_ai_input_size = 2048
|
|
operator.generate_normal = True
|
|
context = SimpleNamespace(
|
|
scene=SimpleNamespace(anyimage_settings=SimpleNamespace())
|
|
)
|
|
with (
|
|
patch.object(self.cutout_main, "require_environment"),
|
|
patch.object(self.cutout_main, "require_model") as require_model,
|
|
patch.object(self.cutout_main, "require_input_path", return_value=Path("input.png")),
|
|
patch.object(
|
|
self.cutout_main, "production_device", return_value="CPU"
|
|
),
|
|
patch.object(
|
|
self.cutout_main,
|
|
"moge2_parameters",
|
|
return_value={"model": operator.model, "resolution_level": 5},
|
|
),
|
|
):
|
|
parameters = operator.request(context)
|
|
|
|
self.assertEqual(
|
|
[call.args[0] for call in require_model.call_args_list],
|
|
[operator.model],
|
|
)
|
|
self.assertTrue(parameters["generate_depth"])
|
|
self.assertEqual(parameters["alpha_threshold"], 0.25)
|
|
self.assertEqual(parameters["normal_mode"], "OBJECT")
|
|
self.assertNotIn("polygon", parameters)
|
|
self.assertNotIn("shape", parameters)
|
|
self.assertNotIn("normal_analysis", parameters)
|
|
|
|
|
|
def test_cutout_geometry_settings_use_configured_moge(self):
|
|
preferences = SimpleNamespace(
|
|
geometry_model="MOGE2_VITB_NORMAL",
|
|
geometry_resolution_level=7,
|
|
)
|
|
with patch.object(self.cutout_main, "addon_preferences", return_value=preferences), patch.object(
|
|
self.anyimage.preferences, "addon_preferences", return_value=preferences
|
|
):
|
|
settings = self.cutout_main.cutout_geometry_settings(None)
|
|
|
|
self.assertEqual(
|
|
settings,
|
|
("MOGE2_VITB_NORMAL", 7),
|
|
)
|
|
|
|
|
|
def test_cutout_tool_settings_gate_normal_generation(self):
|
|
class Row:
|
|
def __init__(self, events):
|
|
self.events = events
|
|
self.enabled = True
|
|
|
|
def prop(self, _owner, name, **_options):
|
|
self.events.append(("prop", name, self.enabled))
|
|
|
|
class Layout:
|
|
def __init__(self):
|
|
self.events = []
|
|
|
|
def operator(self, identifier, **_options):
|
|
self.events.append(("operator", identifier, True))
|
|
|
|
def row(self):
|
|
return Row(self.events)
|
|
|
|
def prop(self, _owner, name, **_options):
|
|
self.events.append(("prop", name, True))
|
|
|
|
context = SimpleNamespace(
|
|
scene=SimpleNamespace(anyimage_settings=SimpleNamespace())
|
|
)
|
|
missing = {
|
|
"environment_ready": False,
|
|
"missing_models": (),
|
|
"ready": False,
|
|
}
|
|
setup = ("operator", "anyimage.setup_ai_environment", True)
|
|
with patch.object(self.ai, "ai_status", return_value=missing):
|
|
cutout_tool = importlib.import_module(
|
|
"anyimage.operators.cutout_tool"
|
|
).CutoutTool
|
|
layout = Layout()
|
|
cutout_tool.draw_settings(context, layout, None)
|
|
self.assertEqual(
|
|
layout.events,
|
|
[
|
|
setup,
|
|
("prop", "cutout_generate_normal", False),
|
|
("prop", "cutout_mesh_detail", True),
|
|
("prop", "cutout_fine_outline", True),
|
|
("prop", "cutout_alpha_threshold", True),
|
|
("prop", "cutout_inward_axis", True),
|
|
],
|
|
)
|
|
|
|
ready = {"environment_ready": True, "missing_models": (), "ready": True}
|
|
with patch.object(self.ai, "ai_status", return_value=ready):
|
|
layout = Layout()
|
|
cutout_tool.draw_settings(context, layout, None)
|
|
self.assertNotIn(setup, layout.events)
|
|
self.assertTrue(all(event[2] for event in layout.events))
|
|
|
|
|
|
def test_cutout_ai_response_reuses_bounds_color_and_lasso(self):
|
|
import numpy as np
|
|
|
|
source = SimpleNamespace(data=SimpleNamespace())
|
|
color_image = SimpleNamespace(users=0)
|
|
pure_values = np.ones((3, 4), dtype=np.float32)
|
|
pure_values[:, 0] = 0.0
|
|
selection_mask = self.image_selection.SelectionMask(
|
|
pure_values,
|
|
(10, 20, 14, 23),
|
|
)
|
|
operator = self.cutout_main.CutoutSelectionToShape()
|
|
operator.alpha_threshold = 0.25
|
|
operator.fine_outline = False
|
|
operator.source_object_name = "Source"
|
|
operator.selection_bounds_json = "[10,20,14,23]"
|
|
operator.input_path = __file__
|
|
operator.color_path = __file__
|
|
operator.shape = "FLAT"
|
|
operator.point_spacing = 32
|
|
operator.generate_normal = True
|
|
operator._selection_mask = selection_mask
|
|
result = SimpleNamespace(value={}, file=lambda _key: Path(__file__))
|
|
|
|
with (
|
|
patch.object(
|
|
self.cutout_main,
|
|
"require_image_empty",
|
|
return_value=source,
|
|
),
|
|
patch.object(
|
|
self.cutout_main,
|
|
"_generated_color_result",
|
|
return_value=(
|
|
color_image,
|
|
np.ones((3, 4), dtype=np.float32),
|
|
selection_mask.bounds,
|
|
),
|
|
) as color_result,
|
|
patch.object(self.cutout_main, "create_cutout_shape") as create_shape,
|
|
):
|
|
message = operator.response(SimpleNamespace(), result)
|
|
|
|
color_result.assert_called_once_with(
|
|
__file__,
|
|
selection_mask.bounds,
|
|
)
|
|
|
|
self.assertEqual(message, "Flat is ready")
|
|
content_values = create_shape.call_args.args[4]
|
|
np.testing.assert_array_equal(content_values, pure_values)
|
|
self.assertEqual(create_shape.call_args.args[5], selection_mask.bounds)
|
|
self.assertIs(create_shape.call_args.kwargs["color_image"], color_image)
|
|
self.assertEqual(create_shape.call_args.kwargs["alpha_threshold"], 0.25)
|
|
self.assertIs(create_shape.call_args.kwargs["fine_outline"], False)
|