nan
08309f719d
- 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.
547 lines
17 KiB
Python
547 lines
17 KiB
Python
import builtins
|
|
import sys
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
|
|
SOURCE_ROOT = Path(__file__).parents[1] / "src"
|
|
sys.path.insert(0, str(SOURCE_ROOT))
|
|
|
|
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.cutout_tool.geometry import balloon as cutout_profile
|
|
from anyimage.operators.cutout_tool.geometry import depth_calibration as cutout_depth
|
|
from anyimage.operators.cutout_tool.geometry.balloon import poisson_balloon_profile
|
|
from anyimage.operators.common.mesh_result import ImageGeometry, load_image_geometry
|
|
from anyimage.operators.common.image_region import (
|
|
RegionMask,
|
|
RegionPath,
|
|
rasterize_region_path,
|
|
)
|
|
from anyimage.operators.convert_to_plane.normal_map import rasterize_tangent_normal
|
|
from anyimage.operators.cutout_tool.geometry.region_contours import (
|
|
RegionContours,
|
|
extract_region_contours,
|
|
)
|
|
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_image_geometry_loader_uses_one_schema_for_plane_and_cutout(tmp_path):
|
|
path = tmp_path / "image_geometry.npz"
|
|
points = np.ones((2, 3, 3), dtype=np.float32)
|
|
valid = np.ones((2, 3), dtype=bool)
|
|
np.savez_compressed(
|
|
path,
|
|
points=points,
|
|
intrinsics=np.eye(3, dtype=np.float32),
|
|
valid=valid,
|
|
calibration_valid=valid,
|
|
image_size=np.asarray((3, 2), dtype=np.int32),
|
|
reference_depth=np.asarray(1.0, dtype=np.float32),
|
|
region_transparent_fraction=np.asarray(0.25, dtype=np.float32),
|
|
)
|
|
|
|
field = load_image_geometry(path)
|
|
|
|
assert field.points.shape == (2, 3, 3)
|
|
assert field.image_size == (3, 2)
|
|
assert field.reference_depth == 1.0
|
|
assert field.region_transparent_fraction == 0.25
|
|
|
|
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 = ImageGeometry(
|
|
points=points,
|
|
intrinsics=np.asarray(
|
|
((100.0, 0.0, 50.0), (0.0, 100.0, 25.0), (0.0, 0.0, 1.0))
|
|
),
|
|
valid=np.ones((51, 101), dtype=bool),
|
|
calibration_valid=np.ones((51, 101), dtype=bool),
|
|
image_size=(101, 51),
|
|
reference_depth=2.0,
|
|
region_transparent_fraction=0.0,
|
|
)
|
|
|
|
uniform_scale = mesh._depth_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.calibration_valid,
|
|
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_region_contours_group_holes_with_their_connected_component():
|
|
alpha = np.zeros((120, 200), dtype=np.uint8)
|
|
alpha[20:100, 50:160] = 255
|
|
alpha[45:75, 80:120] = 0
|
|
alpha[5:15, 5:15] = 255
|
|
|
|
region_contours = extract_region_contours(
|
|
RegionMask(alpha, (0, 0, 200, 120)),
|
|
max_resolution=200,
|
|
)
|
|
|
|
assert region_contours.image_size == (200, 120)
|
|
assert sorted(len(component) for component in region_contours.components) == [1, 2]
|
|
contours = [
|
|
contour
|
|
for component in region_contours.components
|
|
for contour in component
|
|
]
|
|
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_region_contours_accept_blender_float_mask_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
|
|
|
|
region_contours = extract_region_contours(
|
|
RegionMask(alpha, (0, 0, 200, 120)),
|
|
max_resolution=200,
|
|
)
|
|
|
|
assert region_contours.image_size == (200, 120)
|
|
assert len(region_contours.components) == 1
|
|
assert len(region_contours.components[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
|
|
|
|
region_contours = mesh._extract_cutout_contours(
|
|
RegionMask(alpha, (0, 0, 16, 12)),
|
|
64,
|
|
threshold=0.5,
|
|
)
|
|
contour = np.asarray(region_contours.components[0][0])
|
|
|
|
assert region_contours.image_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)
|
|
|
|
with patch("builtins.__import__", side_effect=without_pillow):
|
|
region_mask = rasterize_region_path(
|
|
(48, 32),
|
|
RegionPath(points=((6, 4), (42, 4), (42, 28), (6, 28))),
|
|
antialias=True,
|
|
)
|
|
antialiased = region_mask.full_values((48, 32))
|
|
region_contours = extract_region_contours(
|
|
RegionMask(antialiased, (0, 0, 48, 32)),
|
|
max_resolution=24,
|
|
)
|
|
|
|
assert region_contours.image_size == (48, 32)
|
|
assert region_contours.components
|
|
assert np.any((antialiased > 0.0) & (antialiased < 1.0))
|
|
|
|
def test_blender_cutout_contour_resolution_tracks_mesh_detail():
|
|
expected = RegionContours(
|
|
components=(((((0, 0), (1, 0), (0, 1))),),),
|
|
image_size=(1, 1),
|
|
)
|
|
region_mask = RegionMask(np.ones((2, 2), dtype=np.float32), (0, 0, 2, 2))
|
|
with patch.object(mesh, "extract_region_contours", return_value=expected) as trace:
|
|
assert mesh._extract_cutout_contours(region_mask, 32) == expected
|
|
assert mesh._extract_cutout_contours(region_mask, 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, "image_empty_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_value_scales_the_ui_without_changing_node_area():
|
|
assert np.isclose(
|
|
mesh._cutout_cleanup_value(
|
|
0.25,
|
|
((1.0, 0.0), (0.0, 1.0)),
|
|
),
|
|
6.25,
|
|
)
|
|
assert np.isclose(
|
|
mesh._cutout_cleanup_value(
|
|
0.25,
|
|
((4.0, 0.0), (0.0, 9.0)),
|
|
),
|
|
6.25 / 6.0,
|
|
)
|
|
|
|
|
|
def test_cutout_depth_calibration_uses_five_percent_transparency_threshold():
|
|
assert not cutout_depth.use_adaptive_depth_calibration(0.0499)
|
|
assert cutout_depth.use_adaptive_depth_calibration(0.05)
|
|
|
|
|
|
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(
|
|
RegionContours(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(
|
|
RegionContours(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(
|
|
RegionContours(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(
|
|
RegionContours(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
|