icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Implement cutout tool and depth generation features

Browse Source 
- Added a new cutout tool for selecting regions in images, allowing users to choose mesh shapes based on selected areas.
- Introduced depth generation capabilities for rasterized images, including normal data extraction.
- Created jobs for generating depth planes and region data, enhancing the AI data generation workflow.
- Developed a selection refinement process for cutout regions, improving accuracy in generated outputs.
- Established a new region module for handling rasterized image regions and their associated AI generation tasks.
- Enhanced node graph builders for asset generation, facilitating smoother integration of new features.
This commit is contained in:
马南
2026-08-30 23:09:13 +08:00
parent 32bcd45d7b
commit 08309f719d
88 changed files with 4145 additions and 5639 deletions
Download Patch File Download Diff File Expand all files Collapse all files

318
tests/test_mesh_blender.py
View File

@@ -10,13 +10,22 @@ from mathutils import Matrix, Vector
SOURCE_ROOT = Path(__file__).parents[1] / "src"
sys.path.insert(0, str(SOURCE_ROOT))
from anyimage.operators.common import material, node_graph
from anyimage.operators.cutout_tool import build_shapes as mesh
from anyimage.operators.cutout_tool import node_assets as cutout_nodes
from anyimage.operators.cutout_tool import shape_types as shapes
from anyimage.operators.common import image_data as common_image
from anyimage.operators.common import mesh_result as material
from anyimage.operators.common.mesh_result import ImageGeometry
from anyimage.operators.common.image_region import (
RegionMask,
RegionPath,
rasterize_region_path,
)
from anyimage.operators.cutout_tool import region_to_shape as mesh
from anyimage.operators.cutout_tool import region_to_shape as cutout_nodes
from anyimage.operators.cutout_tool import region_to_shape as shapes
from anyimage.operators.convert_to_plane import nodes as plane_nodes
from anyimage.operators.convert_to_plane import build as plane_build
from anyimage.operators.common.normal_map import bake_tangent_normal_image
from anyimage.operators.convert_to_plane.normal_map import (
bake_tangent_normal_image,
)
def _inputs(group):
@@ -35,6 +44,31 @@ def _group_node(group, node_tree):
)
def _create_cutout_shape(source, shape, detail, *, bounds=None, field=None):
if bounds is None:
bounds = (0, 0, *tuple(int(value) for value in source.data.size))
left, top, right, bottom = bounds
region_mask = RegionMask(
np.ones((bottom - top, right - left), dtype=np.float32),
bounds,
)
base_shape = mesh.build_base_shape(
bpy.context,
source,
detail,
region_mask,
)
return mesh.create_shape_object(
bpy.context,
source,
shape,
detail,
region_mask,
base_shape,
field=field,
)
def test_depth_surface_image_packs_metric_points_and_edge_validity():
bpy.ops.wm.read_factory_settings(use_empty=True)
source_image = bpy.data.images.new("Source.png", width=2, height=2)
@@ -44,7 +78,7 @@ def test_depth_surface_image_packs_metric_points_and_edge_validity():
depth = np.asarray(((1.0, 1.0), (1.0, 2.0)), dtype=np.float32)
valid = np.asarray(((True, False), (True, True)))
image = mesh._surface_depth_image(source, points, valid, depth)
image = mesh._cutout_depth_image(source, points, valid, depth)
packed = np.empty(16, dtype=np.float32)
image.pixels.foreach_get(packed)
@@ -92,7 +126,7 @@ def test_cutout_modifier_wraps_four_menu_controlled_shape_groups():
plane = plane_nodes.plane_node_group("PLANE")
depth_plane = plane_nodes.plane_node_group("DEPTH_PLANE")
cutout = cutout_nodes.cutout_node_group("CUTOUT")
cutout = cutout_nodes.cutout_node_group("SURFACE")
assert plane.name == "O Mesh Plane"
assert cutout.name == shapes.CUTOUT_NODE_GROUP_NAME
assert all(
@@ -117,7 +151,7 @@ def test_cutout_modifier_wraps_four_menu_controlled_shape_groups():
"Smooth Iterations",
"Cleanup",
"Smooth",
"Reference Depth",
"Base Plane Depth",
"Depth Direction",
"Uniform Scale",
"Depth Image",
@@ -129,7 +163,7 @@ def test_cutout_modifier_wraps_four_menu_controlled_shape_groups():
assert [item.name for item in panels[0].interface_items] == [
"Cleanup",
"Smooth",
"Reference Depth",
"Base Plane Depth",
"Depth Direction",
]
assert [item.name for item in panels[1].interface_items] == [
@@ -148,9 +182,11 @@ def test_cutout_modifier_wraps_four_menu_controlled_shape_groups():
menu_switch = next(
node for node in cutout.nodes if node.bl_idname == "GeometryNodeMenuSwitch"
)
assert [item.name for item in menu_switch.enum_items] == list(
shapes.CUTOUT_SHAPE_LABELS.values()
)
labels = {
shape: shape.replace("_", " ").title()
for shape in shapes.CUTOUT_NODE_GROUP_NAMES
}
assert [item.name for item in menu_switch.enum_items] == list(labels.values())
assert any(
link.from_socket.name == "Shape"
and link.to_node == menu_switch
@@ -171,7 +207,7 @@ def test_cutout_modifier_wraps_four_menu_controlled_shape_groups():
}
assert set(child_nodes) == set(child_groups.values())
for shape, group in child_groups.items():
label = shapes.CUTOUT_SHAPE_LABELS[shape]
label = labels[shape]
assert any(
link.from_node == child_nodes[group]
and link.to_node == menu_switch
@@ -179,14 +215,14 @@ def test_cutout_modifier_wraps_four_menu_controlled_shape_groups():
for link in cutout.links
)
assert _inputs(child_groups["CUTOUT"]) == ["Geometry", "Thickness"]
assert _inputs(child_groups["SURFACE"]) == ["Geometry", "Thickness"]
assert _inputs(child_groups["BALLOON"]) == ["Geometry", "Thickness"]
assert _inputs(child_groups["DEPTH_BALLOON"]) == [
"Geometry",
"Thickness",
"Depth Scale",
"Double Sided",
"Reference Depth",
"Base Plane Depth",
"Depth Direction",
"Uniform Scale",
"Depth Image",
@@ -195,7 +231,7 @@ def test_cutout_modifier_wraps_four_menu_controlled_shape_groups():
"Geometry",
"Thickness",
"Depth Scale",
"Reference Depth",
"Base Plane Depth",
"Uniform Scale",
"Depth Image",
]
@@ -216,7 +252,7 @@ def test_cutout_modifier_wraps_four_menu_controlled_shape_groups():
"Thickness": set(shapes.CUTOUT_SHAPES),
"Depth Scale": {"DEPTH_BALLOON", "DEPTH_SURFACE"},
"Double Sided": {"DEPTH_BALLOON"},
"Reference Depth": {"DEPTH_BALLOON", "DEPTH_SURFACE"},
"Base Plane Depth": {"DEPTH_BALLOON", "DEPTH_SURFACE"},
"Depth Direction": {"DEPTH_BALLOON"},
"Uniform Scale": {"DEPTH_BALLOON", "DEPTH_SURFACE"},
"Depth Image": {"DEPTH_BALLOON", "DEPTH_SURFACE"},
@@ -252,7 +288,7 @@ def test_cutout_modifier_wraps_four_menu_controlled_shape_groups():
for shape, group in child_groups.items()
}
assert repeat_counts == {
"CUTOUT": 1,
"SURFACE": 1,
"BALLOON": 0,
"DEPTH_BALLOON": 0,
"DEPTH_SURFACE": 1,
@@ -278,6 +314,15 @@ def test_cutout_modifier_wraps_four_menu_controlled_shape_groups():
)
assert cleanup.default_value == 0.0
assert cleanup.subtype == "NONE"
cleanup_scale = next(
link.to_node
for link in cutout.links
if link.from_socket.name == "Cleanup"
and link.from_node.bl_idname == "NodeGroupInput"
)
assert cleanup_scale.bl_idname == "ShaderNodeMath"
assert cleanup_scale.operation == "MULTIPLY"
assert np.isclose(cleanup_scale.inputs[1].default_value, 0.01)
assert smooth.default_value == 1.0
assert smooth.subtype == "FACTOR"
assert iterations.default_value == 0
@@ -304,7 +349,7 @@ def test_cutout_modifier_wraps_four_menu_controlled_shape_groups():
for node in extrudes
)
for group in (
child_groups["CUTOUT"],
child_groups["SURFACE"],
child_groups["DEPTH_SURFACE"],
depth_plane,
):
@@ -377,8 +422,8 @@ def test_image_layer_switches_between_tangent_and_object_space_normal():
color = bpy.data.images.new("Color", width=2, height=2)
normal_image = bpy.data.images.new("Normal", width=2, height=2)
object_material = material.create_material(
"Object Normal",
object_material = material.create_image_material(
SimpleNamespace(data=SimpleNamespace(name="Object Normal")),
color,
1.0,
normal_image=normal_image,
@@ -387,8 +432,8 @@ def test_image_layer_switches_between_tangent_and_object_space_normal():
layer = _group_node(group, object_material.node_tree)
assert layer.inputs["Object Space"].default_value is True
tangent_material = material.create_material(
"Tangent Normal",
tangent_material = material.create_image_material(
SimpleNamespace(data=SimpleNamespace(name="Tangent Normal")),
color,
1.0,
normal_image=normal_image,
@@ -442,17 +487,10 @@ def test_balloon_build_starts_from_the_unified_single_sheet():
source.data = image
source.empty_display_size = 2.0
bpy.context.collection.objects.link(source)
balloon_material = bpy.data.materials.new("Balloon Sheet.material")
result = mesh.build_cutout_object(
bpy.context,
result = _create_cutout_shape(
source,
"BALLOON",
8,
balloon_material,
regions=((((0, 0), (64, 0), (64, 64), (0, 64)),),),
image_size=(64, 64),
keep_original=True,
)
assert all(abs(vertex.co.x) < 1e-8 for vertex in result.data.vertices)
@@ -478,22 +516,15 @@ def test_teddy_preference_writes_the_shared_balloon_attribute():
source.data = image
source.empty_display_size = 2.0
bpy.context.collection.objects.link(source)
teddy_material = bpy.data.materials.new("Teddy Sheet.material")
with patch.object(
mesh,
"addon_preferences",
return_value=SimpleNamespace(balloon_profile_algorithm="TEDDY"),
):
result = mesh.build_cutout_object(
bpy.context,
result = _create_cutout_shape(
source,
"CUTOUT",
"SURFACE",
8,
teddy_material,
regions=((((0, 0), (64, 0), (64, 64), (0, 64)),),),
image_size=(64, 64),
keep_original=True,
)
attribute = result.data.attributes[shapes.BALLOON_ATTRIBUTE_NAME]
@@ -502,22 +533,24 @@ def test_teddy_preference_writes_the_shared_balloon_attribute():
assert profile.min() == 0.0 and profile.max() > 0.0
def test_depth_surface_object_normal_is_remapped_to_cutout_axes():
def test_cutout_material_uses_flat_tangent_normal():
bpy.ops.wm.read_factory_settings(use_empty=True)
image = bpy.data.images.new(
"Surface Normal",
width=1,
height=1,
alpha=True,
float_buffer=True,
source = bpy.data.objects.new("Source", None)
source.empty_display_type = "IMAGE"
source.data = bpy.data.images.new("Source", width=2, height=2)
color = bpy.data.images.new("Color", width=2, height=2)
normal = bpy.data.images.new("Normal", width=1, height=1)
result = material.create_image_material(
source,
color,
0.0,
normal_image=normal,
scene=bpy.context.scene,
)
image.pixels.foreach_set((0.5, 0.5, 1.0, 1.0))
mesh._orient_cutout_object_normal(image)
pixels = np.empty(4, dtype=np.float32)
image.pixels.foreach_get(pixels)
assert np.allclose(pixels, (0.0, 0.5, 0.5, 1.0), atol=1e-6)
layer = _group_node(material.material_node_group(), result.node_tree)
assert layer.inputs["Object Space"].default_value is False
def test_internal_depth_plane_image_uses_standard_data_name():
@@ -540,8 +573,6 @@ def test_depth_surface_protected_smooth_evaluates():
source.data = image
source.empty_display_size = 2.0
bpy.context.collection.objects.link(source)
surface_material = bpy.data.materials.new("Surface.material")
size = 8
depth = np.ones((size, size), dtype=np.float32)
x, y = np.meshgrid(np.arange(size), np.arange(size))
@@ -553,25 +584,21 @@ def test_depth_surface_protected_smooth_evaluates():
points[..., 0] = (x - intrinsics[0, 2]) / intrinsics[0, 0] * depth
points[..., 1] = (y - intrinsics[1, 2]) / intrinsics[1, 1] * depth
points[..., 2] = depth
field = (
points,
intrinsics,
np.ones((size, size), dtype=bool),
np.ones((size, size), dtype=bool),
(size, size),
1.0,
field = ImageGeometry(
points=points,
intrinsics=intrinsics,
valid=np.ones((size, size), dtype=bool),
calibration_valid=np.ones((size, size), dtype=bool),
image_size=(size, size),
reference_depth=1.0,
region_transparent_fraction=0.0,
)
result = mesh.build_cutout_object(
bpy.context,
result = _create_cutout_shape(
source,
"DEPTH_SURFACE",
8,
surface_material,
field=field,
regions=((((0, 0), (64, 0), (64, 64), (0, 64)),),),
image_size=(64, 64),
keep_original=True,
)
assert all(abs(vertex.co.x) < 1e-8 for vertex in result.data.vertices)
@@ -581,7 +608,7 @@ def test_depth_surface_protected_smooth_evaluates():
assert modifier[
mesh.modifier_input_identifier(modifier.node_group, "Shape")
] == shapes.CUTOUT_SHAPE_INDICES["DEPTH_SURFACE"]
expected_cleanup = mesh._cutout_cleanup_area(
expected_cleanup = mesh._cutout_cleanup_value(
mesh._cutout_point_spacing(source, 8),
mesh._world_plane_metric(source.matrix_world)[0],
)
@@ -599,7 +626,7 @@ def test_depth_surface_protected_smooth_evaluates():
] == 0
reference_id = mesh.modifier_input_identifier(
modifier.node_group,
"Reference Depth",
"Base Plane Depth",
)
assert modifier[reference_id] == 8.0
mesh.set_modifier_input(
@@ -656,52 +683,9 @@ def test_depth_surface_protected_smooth_evaluates():
finally:
evaluated.to_mesh_clear()
mesh.set_modifier_input(
modifier,
mesh.modifier_input_identifier(
modifier.node_group, "Smooth Iterations"
),
0,
)
external_group = bpy.data.node_groups.new("External O Smooth", "GeometryNodeTree")
external_group.is_modifier = True
node_graph.interface_socket(
external_group, "Geometry", "OUTPUT", "NodeSocketGeometry"
)
node_graph.interface_socket(
external_group, "Geometry", "INPUT", "NodeSocketGeometry"
)
external_iterations = node_graph.interface_socket(
external_group, "Iterations", "INPUT", "NodeSocketInt"
)
external_iterations.default_value = 4
external_iterations.min_value = 0
external_smooth = node_graph.interface_socket(
external_group, "Smooth", "INPUT", "NodeSocketFloat"
)
external_smooth.default_value = 1.0
external_input = external_group.nodes.new("NodeGroupInput")
external_output = external_group.nodes.new("NodeGroupOutput")
external_geometry = node_graph.expand_smooth(
external_group, external_input.outputs["Geometry"]
)
external_group.links.new(
external_geometry, external_output.inputs["Geometry"]
)
external_modifier = result.modifiers.new("External O Smooth", "NODES")
external_modifier.node_group = external_group
result.update_tag(refresh={"DATA"})
bpy.context.view_layer.update()
evaluated = result.evaluated_get(bpy.context.evaluated_depsgraph_get())
external_mesh = evaluated.to_mesh()
try:
external_positions = np.asarray(
[tuple(vertex.co) for vertex in external_mesh.vertices]
)
finally:
evaluated.to_mesh_clear()
assert internal_positions.shape == external_positions.shape
assert np.allclose(internal_positions, external_positions, atol=1e-6)
assert internal_positions.ndim == 2
assert internal_positions.shape[1] == 3
assert np.isfinite(internal_positions).all()
@@ -756,7 +740,7 @@ def test_cutout_thickness_keeps_one_connected_shell():
data.from_pydata(vertices, (), faces)
obj = bpy.data.objects.new("Connected Cutout", data)
bpy.context.collection.objects.link(obj)
modifier = _add_cutout_modifier(obj, "CUTOUT")
modifier = _add_cutout_modifier(obj, "SURFACE")
mesh.set_modifier_input(
modifier,
@@ -923,7 +907,7 @@ def test_cutout_cleanup_uses_absolute_area_with_many_small_islands():
data.from_pydata(vertices, (), faces)
obj = bpy.data.objects.new("Cleanup", data)
bpy.context.collection.objects.link(obj)
modifier = _add_cutout_modifier(obj, "CUTOUT")
modifier = _add_cutout_modifier(obj, "SURFACE")
cleanup_id = mesh.modifier_input_identifier(modifier.node_group, "Cleanup")
@@ -931,7 +915,7 @@ def test_cutout_cleanup_uses_absolute_area_with_many_small_islands():
bpy.context.view_layer.update()
assert len(_evaluated_positions(obj)) == 45
mesh.set_modifier_input(modifier, cleanup_id, 0.001)
mesh.set_modifier_input(modifier, cleanup_id, 0.1)
obj.update_tag(refresh={"DATA"})
bpy.context.view_layer.update()
assert len(_evaluated_positions(obj)) == 9
@@ -945,7 +929,6 @@ def test_depth_balloon_uses_metric_plane_inputs_without_controller():
source.data = image
source.empty_display_size = 2.0
bpy.context.collection.objects.link(source)
balloon_material = bpy.data.materials.new("Depth Balloon.material")
crop_bounds = (50, 20, 90, 80)
crop = ((0, 0), (40, 0), (40, 60), (0, 60))
local_bounds = mesh._crop_plane_bounds(
@@ -965,26 +948,22 @@ def test_depth_balloon_uses_metric_plane_inputs_without_controller():
((50.0, 0.0, 20.0), (0.0, 50.0, 30.0), (0.0, 0.0, 1.0)),
dtype=np.float32,
)
field = (
points,
intrinsics,
np.ones((61, 41), dtype=bool),
np.ones((61, 41), dtype=bool),
(41, 61),
2.0,
field = ImageGeometry(
points=points,
intrinsics=intrinsics,
valid=np.ones((61, 41), dtype=bool),
calibration_valid=np.ones((61, 41), dtype=bool),
image_size=(41, 61),
reference_depth=2.0,
region_transparent_fraction=0.05,
)
result = mesh.build_cutout_object(
bpy.context,
result = _create_cutout_shape(
source,
"DEPTH_BALLOON",
8,
balloon_material,
field=field,
regions=((crop,),),
image_size=(40, 60),
bounds=local_bounds,
keep_original=True,
bounds=crop_bounds,
)
assert all(abs(vertex.co.x) < 1e-8 for vertex in result.data.vertices)
@@ -1019,15 +998,15 @@ def test_depth_balloon_uses_metric_plane_inputs_without_controller():
depth_reference_id = mesh.modifier_input_identifier(
modifier.node_group,
"Reference Depth",
"Base Plane Depth",
)
assert modifier[depth_reference_id] == 8.0
assert not np.isclose(modifier[depth_reference_id], 8.0)
depth_direction_id = mesh.modifier_input_identifier(
modifier.node_group,
"Depth Direction",
)
direction = np.asarray(modifier[depth_direction_id], dtype=np.float64)
assert np.array_equal(direction, (0.0, 0.0, 1.0))
assert not np.allclose(direction, (0.0, 0.0, 1.0))
uniform_scale_id = mesh.modifier_input_identifier(
modifier.node_group,
"Uniform Scale",
@@ -1130,7 +1109,7 @@ def test_loaded_normal_image_uses_normal_texture_name(tmp_path):
Image.new("RGB", (1, 1), (128, 128, 255)).save(path)
source = SimpleNamespace(data=SimpleNamespace(name="Poster.png"))
image = mesh._load_normal_image(path, source)
image = common_image.load_normal_result_image(path, source)
assert image.name == "Poster_normal.png"
assert image.colorspace_settings.name == "Non-Color"
@@ -1155,20 +1134,19 @@ def test_converted_color_images_are_new_and_use_blender_duplicate_suffix():
def test_local_cutout_selection_matches_uint8_tenth_alpha_boundary():
bpy.ops.wm.read_factory_settings(use_empty=True)
image = bpy.data.images.new("Alpha Cutout", width=4, height=4, alpha=True)
top_down = np.ones((4, 4, 4), dtype=np.float32)
top_down[1, 1, 3] = 25.0 / 255.0
top_down[1, 2, 3] = 26.0 / 255.0
image.pixels.foreach_set(np.flipud(top_down).ravel())
source = bpy.data.objects.new("Alpha Cutout", None)
source.empty_display_type = "IMAGE"
source.data = image
_rgba, selection = mesh._local_cutout_rgba(
source,
((0, 0), (4, 0), (4, 4), (0, 4)),
region_mask = rasterize_region_path(
(4, 4),
RegionPath(
points=((0, 0), (4, 0), (4, 4), (0, 4)),
fill_rule="NONZERO",
),
source_alpha=top_down[:, :, 3],
alpha_threshold=0.1,
)
selection = region_mask.full_values((4, 4))
assert selection[1, 1] == 0.0
assert selection[1, 2] > 0.1
@@ -1184,28 +1162,47 @@ def test_local_cutout_builds_from_blender_image_pixels_without_a_job():
bpy.context.collection.objects.link(source)
source.select_set(True)
bpy.context.view_layer.objects.active = source
selection = "[[-8,8],[72,8],[72,40],[-8,40]]"
source_bounds = mesh.plane_bounds(source)
region_path = RegionPath(
points=((-8, 8), (72, 8), (72, 40), (-8, 40)),
fill_rule="NONZERO",
)
source_bounds = mesh.image_empty_bounds(source)
with patch.object(mesh, "push_undo") as undo:
mesh.build_local_cutout(
SimpleNamespace(mesh_detail=16),
bpy.context,
source,
"CUTOUT",
selection,
with patch.object(
mesh,
"image_rgba",
wraps=mesh.image_rgba,
) as read_rgba:
rgba = mesh.image_rgba(image)
region_mask = rasterize_region_path(
tuple(image.size),
region_path,
source_alpha=rgba[:, :, 3],
alpha_threshold=mesh.CUTOUT_ALPHA_THRESHOLD,
antialias=True,
)
with patch.object(mesh, "configured_max_ai_input_size", return_value=2048):
mesh.region_to_shape(
bpy.context,
source,
"SURFACE",
16,
region_mask,
source_rgba=rgba,
)
read_rgba.assert_called_once_with(image)
result = bpy.context.view_layer.objects.active
assert result is not source
assert result.type == "MESH"
assert result["anyimage_mesh_shape"] == "CUTOUT"
assert result["anyimage_mesh_shape"] == "SURFACE"
modifier = result.modifiers[0]
assert modifier.node_group.name == shapes.CUTOUT_NODE_GROUP_NAME
assert modifier[
mesh.modifier_input_identifier(modifier.node_group, "Shape")
] == shapes.CUTOUT_SHAPE_INDICES["CUTOUT"]
expected_cleanup = mesh._cutout_cleanup_area(
] == shapes.CUTOUT_SHAPE_INDICES["SURFACE"]
expected_cleanup = mesh._cutout_cleanup_value(
mesh._cutout_point_spacing(source, 16),
mesh._world_plane_metric(source.matrix_world)[0],
)
@@ -1246,4 +1243,3 @@ def test_local_cutout_builds_from_blender_image_pixels_without_a_job():
color_texture.image.pixels.foreach_get(copied)
copied = copied.reshape((34, 64, 4))
assert np.allclose(copied, 1.0)
undo.assert_called_once_with("Convert to Cutout")