Files
AnyImage/tests/test_mesh.py
nan 32bcd45d7b Add new node asset scripts for geometry and shader operations
- Implemented `build_mesh_plane.py` to create a geometry node group for mesh planes with customizable parameters.
- Created `build_shadeless.py` for a shader node group that displays color without scene lighting while preserving shadows and reflections.
- Added `validate.py` to validate the integrity of the newly created node groups and their inputs.
- Introduced `validate_cutout.py` to ensure the functionality of cutout node groups and their dependencies.
- Developed `validate_mesh_plane.py` to test the mesh plane creation and its modifier settings.
2026-08-30 00:17:29 +08:00

562 lines
18 KiB
Python

import builtins
import sys
import tempfile
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
import numpy as np
from PIL import Image
SOURCE_ROOT = Path(__file__).parents[1] / "src"
sys.path.insert(0, str(SOURCE_ROOT))
from anyimage.operators.cutout_tool import balloon_profile as cutout_profile
from anyimage.operators.cutout_tool import build_shapes as mesh
from anyimage.operators.cutout_tool import depth_calibration as cutout_depth
from anyimage.operators.cutout_tool import node_assets as cutout_nodes
from anyimage.operators.cutout_tool import shape_types as shapes
from anyimage.operators.cutout_tool.balloon_profile import poisson_balloon_profile
from anyimage.operators.common.normal_map import rasterize_tangent_normal
from anyimage.operators.cutout_tool.mask_regions import alpha_regions, image_alpha_regions
from anyimage.server.models import moge2 as moge2_model
def test_cutout_mesh_detail_uses_the_cropped_analysis_scale():
assert mesh.limited_cutout_mesh_detail(256, (4096, 1024), 2048) == 128
assert mesh.limited_cutout_mesh_detail(256, (1024, 512), 2048) == 256
assert mesh.limited_cutout_mesh_detail(8, (8192, 4096), 2048) == 4
def test_cutout_node_groups_are_loaded_from_the_asset():
loaded = object()
with patch.object(
cutout_nodes, "load_asset_node_group", return_value=loaded
) as load:
for shape in shapes.CUTOUT_NODE_GROUP_NAMES:
assert cutout_nodes.cutout_node_group(shape) is loaded
load.assert_called_with(shapes.CUTOUT_NODE_GROUP_NAME)
assert load.call_count == 4
def test_base_plane_fit_is_stable_for_a_slender_diagonal_domain():
along = np.linspace(-1.0, 1.0, 101)
across = np.resize(np.asarray((-0.01, 0.0, 0.01)), len(along))
coordinates = np.column_stack(
(
along - across,
along + across,
)
)
values = 0.4 * along + np.sign(across) * 0.05
plane = cutout_depth._stable_plane(coordinates, values)
selected_plane = coordinates @ plane[:2] + plane[2]
corners = np.asarray(
((-1.0, -1.0), (-1.0, 1.0), (1.0, -1.0), (1.0, 1.0))
)
corner_plane = corners @ plane[:2] + plane[2]
assert np.ptp(corner_plane) <= np.ptp(selected_plane) * 1.25
assert np.corrcoef(selected_plane, along)[0, 1] > 0.99
def test_depth_balloon_plane_calibration_matches_cutout_axes():
vertices = []
uv = []
camera_depth = np.empty((3, 3), dtype=np.float32)
for row in range(3):
local_z = 1.0 - row
for column in range(3):
local_y = column - 1.0
vertices.append((0.0, local_y, local_z))
uv.append((column * 0.5, 1.0 - row * 0.5))
camera_depth[row, column] = 2.0 + 0.2 * local_y + 0.1 * local_z
reference, direction = cutout_depth.depth_balloon_projection_parameters(
vertices,
uv,
camera_depth,
np.ones((3, 3), dtype=bool),
0.5,
)
expected_direction = np.asarray((-0.1, 0.05, 1.0), dtype=np.float64)
expected_direction /= np.linalg.norm(expected_direction)
assert np.isclose(
reference,
np.percentile(camera_depth, 95.0) * 0.5,
atol=1e-6,
)
assert np.allclose(direction, expected_direction, atol=1e-6)
def test_depth_balloon_reference_uses_the_shared_upper_depth_percentile():
size = 10
row, column = np.indices((size, size))
camera_depth = np.full((size, size), 2.0, dtype=np.float32)
camera_depth[(row + column) % 2 == 0] = 2.2
vertices = []
uv = []
for y in range(size):
for x in range(size):
vertices.append((0.0, float(x), float(size - 1 - y)))
uv.append((x / (size - 1), 1.0 - y / (size - 1)))
reference, direction = cutout_depth.depth_balloon_projection_parameters(
vertices,
uv,
camera_depth,
np.ones((size, size), dtype=bool),
1.0,
)
assert reference > float(np.median(camera_depth))
assert np.isclose(reference, 2.2, atol=1e-6)
assert np.allclose(direction, (0.0, 0.0, 1.0), atol=1e-6)
def mesh_component_count(faces):
adjacency = {}
for face in faces:
for vertex in face:
adjacency.setdefault(int(vertex), set())
for start, end in zip(face, np.roll(face, -1)):
adjacency[int(start)].add(int(end))
adjacency[int(end)].add(int(start))
remaining = set(adjacency)
count = 0
while remaining:
count += 1
pending = [remaining.pop()]
while pending:
vertex = pending.pop()
connected = adjacency[vertex] & remaining
remaining.difference_update(connected)
pending.extend(connected)
return count
def test_tangent_normal_pixels_use_the_final_geometry_frame():
object_normal = np.zeros((3, 3, 3), dtype=np.float32)
object_normal[..., 2] = 1.0
triangle_uv = np.asarray((((0.0, 0.0), (1.0, 0.0), (0.0, 1.0)),))
triangle_normal = np.asarray((((1.0, 0.0, 0.0),) * 3,))
triangle_tangent = np.asarray((((0.0, 1.0, 0.0),) * 3,))
triangle_sign = np.ones((1, 3), dtype=np.float32)
tangent = rasterize_tangent_normal(
object_normal,
triangle_uv,
triangle_normal,
triangle_tangent,
triangle_sign,
)
assert np.allclose(tangent[0, 0], (0.0, 1.0, 0.0))
assert not np.allclose(tangent[0, 0], object_normal[0, 0])
assert np.allclose(tangent[2, 2], (0.0, 0.0, 1.0))
def test_depth_surface_discontinuity_map_uses_local_three_by_three_range():
depth = np.ones((5, 5), dtype=np.float32)
depth[:, 3:] = 2.0
valid = np.ones_like(depth, dtype=bool)
discontinuity = mesh._depth_discontinuity_map(depth, valid)
assert discontinuity[2, 0] == 0.0
assert discontinuity[2, 2] == 1.0
def test_depth_surface_projection_uses_one_scale_for_all_camera_axes():
points = np.zeros((51, 101, 3), dtype=np.float32)
points[..., 2] = np.linspace(1.0, 3.0, 51 * 101).reshape(51, 101)
field = (
points,
np.asarray(((100.0, 0.0, 50.0), (0.0, 100.0, 25.0), (0.0, 0.0, 1.0))),
np.ones((51, 101), dtype=bool),
np.ones((51, 101), dtype=bool),
(101, 51),
2.0,
)
uniform_scale = mesh._surface_uniform_scale(
SimpleNamespace(),
field,
bounds=(-2.0, 2.0, -1.0, 1.0),
)
assert uniform_scale == 2.0
reference_depth = cutout_depth.metric_depth_reference(
points[..., 2],
field[3],
uniform_scale,
)
assert np.isclose(reference_depth, 5.8, atol=1e-6)
def test_moge_surface_field_uses_the_model_validity_mask():
prediction = moge2_model.Prediction(
depth=np.ones((1, 2, 2), dtype=np.float32),
normal=np.zeros((1, 2, 2, 3), dtype=np.float32),
mask=np.asarray((((True, False), (True, True)),)),
intrinsics=np.asarray(
(((1.0, 0.0, 0.5), (0.0, 1.0, 0.5), (0.0, 0.0, 1.0)),),
dtype=np.float32,
),
points=np.ones((1, 2, 2, 3), dtype=np.float32),
)
frame = moge2_model.geometry_frame(prediction)
assert np.array_equal(frame.valid_mask, prediction.mask[0])
def test_cutout_regions_group_holes_with_their_connected_component():
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "cutout.png"
pixels = np.zeros((120, 200, 4), dtype=np.uint8)
pixels[20:100, 50:160] = (20, 40, 60, 255)
pixels[45:75, 80:120] = 0
pixels[5:15, 5:15] = (20, 40, 60, 255)
Image.fromarray(pixels, mode="RGBA").save(path)
regions, size = image_alpha_regions(path, max_resolution=200)
assert size == (200, 120)
assert sorted(len(region) for region in regions) == [1, 2]
contours = [contour for region in regions for contour in region]
bounds = [
(*np.min(contour, axis=0), *np.max(contour, axis=0))
for contour in contours
]
assert (50.0, 20.0, 160.0, 100.0) in bounds
assert (80.0, 45.0, 120.0, 75.0) in bounds
assert (5.0, 5.0, 15.0, 15.0) in bounds
def test_alpha_regions_accepts_blender_float_alpha_with_a_hole():
alpha = np.zeros((120, 200), dtype=np.float32)
alpha[20:100, 50:160] = 1.0
alpha[45:75, 80:120] = 0.0
regions, size = alpha_regions(alpha, max_resolution=200)
assert size == (200, 120)
assert len(regions) == 1
assert len(regions[0]) == 2
def test_cutout_crop_bounds_map_to_the_source_plane():
source = SimpleNamespace(
data=SimpleNamespace(size=(200, 100)),
empty_display_size=2.0,
empty_image_offset=(-0.5, -0.5),
)
bounds = mesh._crop_plane_bounds(
source,
(50, 25, 150, 75),
(200, 100),
)
np.testing.assert_allclose(bounds, (-0.5, 0.5, -0.25, 0.25))
def test_depth_surface_cutout_uses_half_alpha_boundary():
alpha = np.zeros((12, 16), dtype=np.float32)
alpha[2:10, 2:14] = 0.49
alpha[4:8, 5:11] = 0.51
regions, size = mesh._cutout_regions(alpha, 64, threshold=0.5)
contour = np.asarray(regions[0][0])
assert size == (16, 12)
np.testing.assert_allclose(contour.min(axis=0), (5.0, 4.0))
np.testing.assert_allclose(contour.max(axis=0), (11.0, 8.0))
def test_blender_cutout_alpha_pipeline_does_not_import_pillow():
original_import = builtins.__import__
def without_pillow(name, *args, **kwargs):
if name == "PIL" or name.startswith("PIL."):
raise ModuleNotFoundError("Pillow is unavailable in Blender")
return original_import(name, *args, **kwargs)
selection = np.zeros((32, 48), dtype=bool)
selection[4:28, 6:42] = True
with patch("builtins.__import__", side_effect=without_pillow):
antialiased = mesh._antialias_selection(selection)
regions, size = alpha_regions(antialiased, max_resolution=24)
assert size == (48, 32)
assert regions
assert np.any((antialiased > 0.0) & (antialiased < 1.0))
def test_unrefined_cutout_geometry_uses_strict_tenth_source_alpha_threshold():
alpha = np.full((40, 60), 0.1, dtype=np.float32)
alpha[10:30, 20:40] = 0.11
polygon = ((0, 0), (60, 0), (60, 40), (0, 40))
regions, size = mesh._geometry_cutout_regions(alpha, polygon, 64)
contour = np.asarray(regions[0][0])
assert size == (60, 40)
np.testing.assert_allclose(contour.min(axis=0), (20.0, 10.0))
np.testing.assert_allclose(contour.max(axis=0), (40.0, 30.0))
def test_cutout_geometry_excludes_zero_source_alpha_without_refine():
alpha = np.zeros((80, 120), dtype=np.float32)
alpha[10:70, 10:110] = 1.0
alpha[30:50, 45:75] = 0.0
polygon = ((10, 10), (110, 10), (110, 70), (10, 70))
regions, size = mesh._geometry_cutout_regions(alpha, polygon, 64)
assert size == (120, 80)
assert len(regions) == 1
assert len(regions[0]) == 2
def test_cutout_geometry_uses_final_alpha_with_refine():
alpha = np.zeros((80, 120), dtype=np.float32)
alpha[10:70, 10:110] = 1.0
alpha[30:50, 45:75] = 0.0
polygon = ((10, 10), (110, 10), (110, 70), (10, 70))
regions, size = mesh._geometry_cutout_regions(
alpha, polygon, 64, refine_selection=True
)
assert size == (120, 80)
assert len(regions) == 1
assert len(regions[0]) == 2
def test_unrefined_cutout_geometry_discards_out_of_frame_selection():
alpha = np.zeros((60, 100), dtype=np.float32)
alpha[10:50] = 1.0
polygon = ((-20, 10), (120, 10), (120, 50), (-20, 50))
regions, size = mesh._geometry_cutout_regions(alpha, polygon, 64)
points = np.concatenate(
[np.asarray(contour) for region in regions for contour in region]
)
assert size == (100, 60)
assert points[:, 0].min() >= 0.0
assert points[:, 0].max() <= 100.0
def test_refined_cutout_geometry_discards_out_of_frame_selection():
alpha = np.zeros((60, 100), dtype=np.float32)
alpha[10:50] = 1.0
polygon = ((-20, 10), (120, 10), (120, 50), (-20, 50))
regions, size = mesh._geometry_cutout_regions(
alpha, polygon, 64, refine_selection=True
)
points = np.concatenate(
[np.asarray(contour) for region in regions for contour in region]
)
assert size == (100, 60)
assert points[:, 0].min() >= 0.0
assert points[:, 0].max() <= 100.0
def test_blender_cutout_contour_resolution_tracks_mesh_detail():
expected = ([(((0, 0), (1, 0), (0, 1)),)], (1, 1))
alpha = np.ones((2, 2), dtype=np.float32)
with patch.object(mesh, "alpha_regions", return_value=expected) as trace:
assert mesh._cutout_regions(alpha, 32) == expected
assert mesh._cutout_regions(alpha, 256) == expected
assert trace.call_args_list[0].kwargs["max_resolution"] == 512
assert trace.call_args_list[1].kwargs["max_resolution"] == 1024
def test_cutout_point_spacing_uses_the_complete_source_plane():
source = SimpleNamespace(matrix_world=object())
with (
patch.object(mesh, "plane_bounds", return_value=(-2.0, 2.0, -1.0, 1.0)),
patch.object(
mesh,
"_world_plane_metric",
return_value=(((1.0, 0.0), (0.0, 1.0)), 1.0),
),
):
spacing = mesh._cutout_point_spacing(source, 200)
assert spacing == 0.02
def test_cutout_cleanup_area_uses_point_spacing_and_plane_area_scale():
assert np.isclose(
mesh._cutout_cleanup_area(
0.25,
((1.0, 0.0), (0.0, 1.0)),
),
0.0625,
)
assert np.isclose(
mesh._cutout_cleanup_area(
0.25,
((4.0, 0.0), (0.0, 9.0)),
),
0.0625 / 6.0,
)
def test_poisson_balloon_profile_is_zero_on_boundary_and_positive_inside():
points = np.asarray(
((0, 0), (2, 0), (2, 2), (0, 2), (1, 1)),
dtype=float,
)
faces = ((0, 1, 4), (1, 2, 4), (2, 3, 4), (3, 0, 4))
boundary = np.asarray((True, True, True, True, False))
profile = poisson_balloon_profile(points, faces, boundary)
assert np.all(profile[:4] == 0.0)
assert profile[4] > 0.0
def test_poisson_balloon_profile_scales_in_world_units():
points = np.asarray(
((0, 0), (2, 0), (2, 2), (0, 2), (1, 1)),
dtype=float,
)
faces = ((0, 1, 4), (1, 2, 4), (2, 3, 4), (3, 0, 4))
boundary = np.asarray((True, True, True, True, False))
original = poisson_balloon_profile(points, faces, boundary)
scaled = poisson_balloon_profile(points * 2.0, faces, boundary)
assert np.isclose(scaled[4], original[4] * 2.0, atol=1e-6)
def test_poisson_balloon_profile_makes_wide_regions_thicker():
faces = ((0, 1, 4), (1, 2, 4), (2, 3, 4), (3, 0, 4))
boundary = np.asarray((True, True, True, True, False))
wide = np.asarray(
((0, 0), (4, 0), (4, 2), (0, 2), (2, 1)),
dtype=float,
)
narrow = np.asarray(
((0, 0), (4, 0), (4, 0.6), (0, 0.6), (2, 0.3)),
dtype=float,
)
wide_profile = poisson_balloon_profile(wide, faces, boundary)
narrow_profile = poisson_balloon_profile(narrow, faces, boundary)
assert wide_profile[4] > narrow_profile[4]
def test_teddy_balloon_profile_uses_chordal_axis_radius():
points = np.asarray(
((0, 0), (4, 0), (4, 2), (0, 2), (2, 1)),
dtype=float,
)
faces = ((0, 1, 4), (1, 2, 4), (2, 3, 4), (3, 0, 4))
boundary = np.asarray((True, True, True, True, False))
nodes = np.asarray(((1, 1), (3, 1)), dtype=float)
with patch.object(
cutout_profile,
"chordal_axis",
return_value=(nodes, [(0, 1)], np.asarray((1.0, 1.0))),
):
profile = cutout_profile.teddy_balloon_heights(
points,
faces,
points[:4],
boundary,
0.45,
)
assert np.all(profile[:4] == 0.0)
assert np.isclose(profile[4], 0.45, atol=1e-6)
def test_cutout_domain_adapts_boundary_spacing_for_a_small_valid_region():
regions = (
(((2046, 2046), (2050, 2046), (2050, 2050), (2046, 2050)),),
)
domain = mesh._cutout_domain(
regions,
(4096, 4096),
(-1.0, 1.0, -1.0, 1.0),
0.5,
((1.0, 0.0), (0.0, 1.0)),
)
assert len(domain[0]) >= 3
assert len(domain[1]) >= 1
def test_cutout_domain_builds_a_full_image_at_minimum_detail():
regions = ((((0, 0), (4096, 0), (4096, 4096), (0, 4096)),),)
domain = mesh._cutout_domain(
regions,
(4096, 4096),
(-1.0, 1.0, -1.0, 1.0),
0.5,
((1.0, 0.0), (0.0, 1.0)),
)
assert len(domain[0]) >= 4
assert len(domain[1]) >= 2
def test_cutout_domain_centers_vertices_on_crop_bounds():
regions = ((((0, 0), (100, 0), (100, 50), (0, 50)),),)
vertices, _faces, uv, _boundary, _planar, _profile = mesh._cutout_domain(
regions,
(100, 50),
(2.0, 6.0, -3.0, 1.0),
0.5,
((1.0, 0.0), (0.0, 1.0)),
)
center = (vertices[:, :2].min(axis=0) + vertices[:, :2].max(axis=0)) * 0.5
assert np.allclose(center, 0.0, atol=1e-6)
expected_uv = vertices[:, :2] / np.asarray((4.0, 4.0)) + 0.5
assert np.allclose(uv, expected_uv, atol=1e-6)
def test_cutout_domain_keeps_overlapping_regions_disconnected():
regions = (
(
((20, 0), (100, 0), (100, 80), (20, 80)),
((50, 20), (50, 60), (70, 60), (70, 20)),
),
(
((0, 20), (40, 20), (40, 50), (0, 50)),
),
)
domain = mesh._cutout_domain(
regions,
(120, 80),
(0.0, 1.0, 0.0, 1.0),
0.025,
((1.0, 0.0), (0.0, 1.0)),
)
centers = domain[2][domain[1]].mean(axis=1)
in_hole = (
(centers[:, 0] > 0.45)
& (centers[:, 0] < 0.55)
& (centers[:, 1] > 0.35)
& (centers[:, 1] < 0.65)
)
in_left_region = (
(centers[:, 0] > 0.02)
& (centers[:, 0] < 0.12)
& (centers[:, 1] > 0.4)
& (centers[:, 1] < 0.75)
)
assert not in_hole.any()
assert in_left_region.any()
assert mesh_component_count(domain[1]) == 2