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

251
tests/test_mesh.py
View File

@@ -1,24 +1,31 @@
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.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
@@ -27,6 +34,29 @@ def test_cutout_mesh_detail_uses_the_cropped_analysis_scale():
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(
@@ -171,16 +201,19 @@ def test_depth_surface_discontinuity_map_uses_local_three_by_three_range():
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,
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._surface_uniform_scale(
uniform_scale = mesh._depth_uniform_scale(
SimpleNamespace(),
field,
bounds=(-2.0, 2.0, -1.0, 1.0),
@@ -189,7 +222,7 @@ def test_depth_surface_projection_uses_one_scale_for_all_camera_axes():
assert uniform_scale == 2.0
reference_depth = cutout_depth.metric_depth_reference(
points[..., 2],
field[3],
field.calibration_valid,
uniform_scale,
)
assert np.isclose(reference_depth, 5.8, atol=1e-6)
@@ -212,20 +245,24 @@ def test_moge_surface_field_uses_the_model_validity_mask():
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)
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
regions, size = image_alpha_regions(path, max_resolution=200)
region_contours = extract_region_contours(
RegionMask(alpha, (0, 0, 200, 120)),
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]
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
@@ -235,16 +272,19 @@ def test_cutout_regions_group_holes_with_their_connected_component():
assert (5.0, 5.0, 15.0, 15.0) in bounds
def test_alpha_regions_accepts_blender_float_alpha_with_a_hole():
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
regions, size = alpha_regions(alpha, max_resolution=200)
region_contours = extract_region_contours(
RegionMask(alpha, (0, 0, 200, 120)),
max_resolution=200,
)
assert size == (200, 120)
assert len(regions) == 1
assert len(regions[0]) == 2
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():
@@ -268,10 +308,14 @@ def test_depth_surface_cutout_uses_half_alpha_boundary():
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])
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 size == (16, 12)
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))
@@ -283,91 +327,31 @@ def test_blender_cutout_alpha_pipeline_does_not_import_pillow():
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)
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 size == (48, 32)
assert regions
assert region_contours.image_size == (48, 32)
assert region_contours.components
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
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
@@ -375,7 +359,7 @@ def test_blender_cutout_contour_resolution_tracks_mesh_detail():
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, "image_empty_bounds", return_value=(-2.0, 2.0, -1.0, 1.0)),
patch.object(
mesh,
"_world_plane_metric",
@@ -387,23 +371,28 @@ def test_cutout_point_spacing_uses_the_complete_source_plane():
assert spacing == 0.02
def test_cutout_cleanup_area_uses_point_spacing_and_plane_area_scale():
def test_cutout_cleanup_value_scales_the_ui_without_changing_node_area():
assert np.isclose(
mesh._cutout_cleanup_area(
mesh._cutout_cleanup_value(
0.25,
((1.0, 0.0), (0.0, 1.0)),
),
0.0625,
6.25,
)
assert np.isclose(
mesh._cutout_cleanup_area(
mesh._cutout_cleanup_value(
0.25,
((4.0, 0.0), (0.0, 9.0)),
),
0.0625 / 6.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)),
@@ -482,8 +471,7 @@ def test_cutout_domain_adapts_boundary_spacing_for_a_small_valid_region():
)
domain = mesh._cutout_domain(
regions,
(4096, 4096),
RegionContours(regions, (4096, 4096)),
(-1.0, 1.0, -1.0, 1.0),
0.5,
((1.0, 0.0), (0.0, 1.0)),
@@ -497,8 +485,7 @@ 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),
RegionContours(regions, (4096, 4096)),
(-1.0, 1.0, -1.0, 1.0),
0.5,
((1.0, 0.0), (0.0, 1.0)),
@@ -511,8 +498,7 @@ 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),
RegionContours(regions, (100, 50)),
(2.0, 6.0, -3.0, 1.0),
0.5,
((1.0, 0.0), (0.0, 1.0)),
@@ -535,8 +521,7 @@ def test_cutout_domain_keeps_overlapping_regions_disconnected():
),
)
domain = mesh._cutout_domain(
regions,
(120, 80),
RegionContours(regions, (120, 80)),
(0.0, 1.0, 0.0, 1.0),
0.025,
((1.0, 0.0), (0.0, 1.0)),