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

860
tests/test_addon.py
View File

File diff suppressed because it is too large Load Diff

51
tests/test_cutout.py
View File

@@ -8,8 +8,9 @@ import bpy
SOURCE_ROOT = Path(__file__).parents[1] / "src"
sys.path.insert(0, str(SOURCE_ROOT))
from anyimage.operators.common.image_gesture import ImageMaskGesture
from anyimage.operators.cutout_tool import selection_tool as cutout
from anyimage.operators.common.image_interaction import ImageRegionGesture
from anyimage.operators.cutout_tool import select_region as cutout
from anyimage.operators.common.image_region import RegionPath
class Pie:
@@ -31,10 +32,10 @@ class WindowManager:
draw(menu, None)
def cutout_context(window_manager, point_spacing=32, normal_analysis=False):
def cutout_context(window_manager, point_spacing=32, normal_mode="NONE"):
settings = SimpleNamespace(
cutout_point_spacing=point_spacing,
cutout_normal_analysis=normal_analysis,
cutout_normal_mode=normal_mode,
)
return SimpleNamespace(
window_manager=window_manager,
@@ -51,10 +52,15 @@ def test_cutout_pie_uses_four_semantic_shapes_in_cardinal_order():
context = cutout_context(window_manager)
source = cutout_source()
cutout._open_shape_pie(context, object(), source, ((0, 0), (10, 0), (0, 10)))
cutout._open_shape_pie(
context,
object(),
source,
RegionPath(points=((0, 0), (10, 0), (0, 10))),
)
assert [call[2].shape for call in window_manager.pie.calls] == [
"CUTOUT",
"SURFACE",
"BALLOON",
"DEPTH_BALLOON",
"DEPTH_SURFACE",
@@ -66,18 +72,18 @@ def test_cutout_pie_uses_four_semantic_shapes_in_cardinal_order():
"Depth Surface",
]
assert all(
call[0] == "anyimage.convert_cutout_shape" for call in window_manager.pie.calls
call[0] == "anyimage.cutout_region_to_shape"
for call in window_manager.pie.calls
)
assert not any(call[2].refine_selection for call in window_manager.pie.calls)
assert not any(call[2].normal_analysis for call in window_manager.pie.calls)
assert not any(call[2].invert for call in window_manager.pie.calls)
assert all(call[2].normal_mode == "NONE" for call in window_manager.pie.calls)
assert all(call[2].mesh_detail == 64 for call in window_manager.pie.calls)
def test_cutout_pie_passes_image_lasso_options_to_every_shape():
window_manager = WindowManager()
context = cutout_context(
window_manager, point_spacing=16, normal_analysis=True
window_manager, point_spacing=16, normal_mode="OBJECT"
)
source = cutout_source(size=(4096, 2048))
@@ -85,14 +91,19 @@ def test_cutout_pie_passes_image_lasso_options_to_every_shape():
context,
object(),
source,
((0, 0), (10, 0), (0, 10)),
RegionPath(
points=((0, 0), (10, 0), (0, 10)),
invert=True,
),
refine_selection=True,
invert=True,
)
assert all(call[2].refine_selection for call in window_manager.pie.calls)
assert all(call[2].normal_analysis for call in window_manager.pie.calls)
assert all(call[2].invert for call in window_manager.pie.calls)
assert all(call[2].normal_mode == "OBJECT" for call in window_manager.pie.calls)
assert all(
RegionPath.from_json(call[2].region_path_json).invert
for call in window_manager.pie.calls
)
assert all(call[2].mesh_detail == 256 for call in window_manager.pie.calls)
@@ -116,8 +127,8 @@ def test_cutout_reuses_image_lasso_gesture_and_double_click_selects_all(monkeypa
cutout, "resolve_image_edit_click", lambda *_args: (source, True)
)
monkeypatch.setattr(cutout, "video_total_frames", lambda _image: 1)
monkeypatch.setattr(cutout, "apply_image_mask_settings", lambda *_args: None)
monkeypatch.setattr(cutout, "selected_image_rgba", lambda *_args: object())
monkeypatch.setattr(cutout, "apply_region_settings", lambda *_args: None)
monkeypatch.setattr(cutout, "_validate_cutout_region", lambda *_args: None)
operator = SimpleNamespace(
full_image=True,
@@ -129,10 +140,12 @@ def test_cutout_reuses_image_lasso_gesture_and_double_click_selects_all(monkeypa
operator, context, SimpleNamespace(value="DOUBLE_CLICK")
)
assert issubclass(cutout.SelectCutoutRegion, ImageMaskGesture)
assert issubclass(cutout.SelectCutoutRegion, ImageRegionGesture)
assert result == {"FINISHED"}
selection = json.loads(window_manager.pie.calls[0][2].selection_data)
assert selection == [[0, 0], [320, 0], [320, 180], [0, 180]]
region_path = RegionPath.from_json(
window_manager.pie.calls[0][2].region_path_json
)
assert region_path.points == ((0, 0), (320, 0), (320, 180), (0, 180))
assert window_manager.pie.calls[0][2].mesh_detail == 10
assert not hasattr(cutout, "pending_selection")

55
tests/test_image.py
View File

@@ -27,21 +27,31 @@ def load_image_modules():
module_paths = (
(
"anyimage.operators.common.image",
PROJECT_ROOT / "src" / "anyimage" / "operators" / "common" / "image.py",
),
(
"anyimage.operators.common.node_asset",
"anyimage.operators.common.image_data",
PROJECT_ROOT
/ "src"
/ "anyimage"
/ "operators"
/ "common"
/ "node_asset.py",
/ "image_data.py",
),
(
"anyimage.operators.common.material",
PROJECT_ROOT / "src" / "anyimage" / "operators" / "common" / "material.py",
"anyimage.operators.common.mesh_result",
PROJECT_ROOT
/ "src"
/ "anyimage"
/ "operators"
/ "common"
/ "mesh_result.py",
),
(
"anyimage.operators.convert_to_plane.resolution",
PROJECT_ROOT
/ "src"
/ "anyimage"
/ "operators"
/ "convert_to_plane"
/ "resolution.py",
),
)
packages = {
@@ -50,6 +60,7 @@ def load_image_modules():
"anyimage",
"anyimage.operators",
"anyimage.operators.common",
"anyimage.operators.convert_to_plane",
)
}
for package in packages.values():
@@ -452,21 +463,7 @@ class ImageTest(unittest.TestCase):
"img_normal.png",
)
def test_copy_depth_image_renames_still_source(self):
class FakeImage:
def __init__(self, name):
self.name = name
def copy(self):
return FakeImage("img.png.001")
source_object = SimpleNamespace(data=FakeImage("img.png"))
depth_image = self.image_depth.copy_depth_image(source_object)
self.assertEqual(depth_image.name, "img_depth.exr")
def test_color_and_material_names_use_source_image_base_name(self):
def test_color_name_uses_source_image_base_name(self):
source_object = SimpleNamespace(
data=SimpleNamespace(name="img.png"),
)
@@ -475,19 +472,11 @@ class ImageTest(unittest.TestCase):
self.image_depth.color_data_name(source_object),
"img_color.png",
)
self.assertEqual(
self.image_depth.material_data_name(source_object),
"img",
)
source_object.data.name = "img.png.001"
self.assertEqual(
self.image_depth.color_data_name(source_object),
"img.001_color.png",
)
self.assertEqual(
self.image_depth.material_data_name(source_object),
"img.001",
)
def test_surface_displacement_is_set_to_displacement_only(self):
class Material:
@@ -501,14 +490,14 @@ class ImageTest(unittest.TestCase):
"DISPLACEMENT",
)
def test_plane_bounds_keep_image_aspect_and_empty_offset(self):
def test_image_empty_bounds_keep_image_aspect_and_empty_offset(self):
source = SimpleNamespace(
data=SimpleNamespace(size=(800, 400)),
empty_display_size=2.0,
empty_image_offset=(-0.5, -0.5),
)
bounds = self.image_depth.plane_bounds(source)
bounds = self.image_depth.image_empty_bounds(source)
self.assertEqual(bounds, (-1.0, 1.0, -0.5, 0.5))

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)),

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")

42
tests/test_pack.py
View File

@@ -67,26 +67,21 @@ class PackTest(unittest.TestCase):
self.assertNotIn("assets/O_AnyImage.blend1", names)
for path in (
"operators/common/__init__.py",
"operators/common/image.py",
"operators/common/material.py",
"operators/common/node_asset.py",
"operators/common/normal_map.py",
"operators/common/operator.py",
"operators/common/image_edit.py",
"operators/common/image_gesture.py",
"operators/common/ai_job.py",
"operators/common/image_data.py",
"operators/common/image_interaction.py",
"operators/common/image_mask.py",
"operators/common/node_graph.py",
"operators/common/image_region.py",
"operators/common/mesh_result.py",
"operators/convert_to_plane/normal_map.py",
"operators/convert_to_plane/resolution.py",
"operators/cutout_tool/__init__.py",
"operators/cutout_tool/balloon_profile.py",
"operators/cutout_tool/build_shapes.py",
"operators/cutout_tool/depth_calibration.py",
"operators/cutout_tool/mask_regions.py",
"operators/cutout_tool/node_assets.py",
"operators/cutout_tool/planar_mesh.py",
"operators/cutout_tool/selection_tool.py",
"operators/cutout_tool/shape_conversion.py",
"operators/cutout_tool/shape_types.py",
"operators/cutout_tool/select_region.py",
"operators/cutout_tool/region_to_shape.py",
"operators/cutout_tool/geometry/__init__.py",
"operators/cutout_tool/geometry/balloon.py",
"operators/cutout_tool/geometry/depth_calibration.py",
"operators/cutout_tool/geometry/region_contours.py",
"operators/cutout_tool/geometry/triangulation.py",
"operators/convert_to_plane/__init__.py",
"operators/convert_to_plane/build.py",
"operators/convert_to_plane/conversion.py",
@@ -118,10 +113,15 @@ class PackTest(unittest.TestCase):
self.assertIn("server/models/geometry.py", names)
self.assertIn("server/jobs/remove_background.py", names)
self.assertIn("server/jobs/refine_image_mask.py", names)
self.assertIn("server/selection/cutout_images.py", names)
self.assertIn("server/selection/polygon_mask.py", names)
self.assertIn("server/jobs/prepare_geometry_data.py", names)
self.assertIn("server/region/generate_selection.py", names)
self.assertIn("server/jobs/generate_region_data.py", names)
self.assertIn("server/jobs/generate_depth_plane.py", names)
self.assertIn("server/geometry/generate_depth.py", names)
self.assertIn("server/geometry/normal_texture.py", names)
self.assertNotIn("operators/common/image_selection.py", names)
self.assertNotIn("server/selection/cutout_images.py", names)
self.assertNotIn("server/selection/polygon_mask.py", names)
self.assertNotIn("server/jobs/prepare_geometry_data.py", names)
self.assertNotIn("server/geometry/cutout.py", names)
self.assertIn("server/models/ben2.py", names)
self.assertIn("server/models/da3.py", names)

701
tests/test_server.py
View File

@@ -41,8 +41,8 @@ class ServerTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.server_module = load_server_module()
importlib.import_module("server.selection.polygon_mask")
importlib.import_module("server.jobs.prepare_geometry_data")
importlib.import_module("server.geometry.generate_depth")
importlib.import_module("server.jobs.generate_region_data")
from blendjob import JobServer
from server.model_manager import ModelManager, ModelSessionCache
@@ -261,7 +261,8 @@ class ServerTest(unittest.TestCase):
"refine-image-mask",
"remove-background",
"upscale-image",
"prepare-geometry-data",
"generate-depth-plane",
"generate-region-data",
"debug-da3",
"debug-moge2",
"debug-ben2",
@@ -385,652 +386,116 @@ class ServerTest(unittest.TestCase):
)
self.assertEqual(events, ["clear", "close"])
def test_polygon_selection_mask_uses_even_odd_fill_for_self_intersection(self):
image_mask = sys.modules["server.selection.polygon_mask"]
polygon = (
(1, 1),
(9, 1),
(9, 9),
(1, 9),
(1, 1),
(3, 3),
(3, 7),
(7, 7),
(7, 3),
(3, 3),
(1, 1),
)
mask = image_mask.polygon_selection_mask((10, 10), polygon)
self.assertTrue(mask[2, 2])
self.assertFalse(mask[4, 4])
self.assertEqual(int(mask.sum()), 48)
def test_polygon_selection_mask_nonzero_keeps_same_direction_overlap(self):
image_mask = sys.modules["server.selection.polygon_mask"]
polygon = (
(1, 1),
(7, 1),
(7, 7),
(1, 7),
(1, 1),
(4, 1),
(10, 1),
(10, 7),
(4, 7),
(4, 1),
(1, 1),
)
even_odd = image_mask.polygon_selection_mask((12, 9), polygon)
nonzero = image_mask.polygon_selection_mask(
(12, 9), polygon, "NONZERO"
)
self.assertTrue(even_odd[3, 2])
self.assertFalse(even_odd[3, 5])
self.assertTrue(nonzero[3, 2])
self.assertTrue(nonzero[3, 5])
self.assertTrue(nonzero[3, 8])
def test_image_mask_canvas_clears_rgba_outside_the_selection(self):
def test_region_transparent_fraction_uses_visible_alpha_threshold(self):
import numpy as np
from PIL import Image
image_mask = sys.modules["server.selection.polygon_mask"]
source_path = self.storage_root / "source.png"
pixels = np.full((2, 2, 4), (20, 40, 60, 255), dtype=np.uint8)
Image.fromarray(pixels, mode="RGBA").save(source_path)
depth = sys.modules["server.geometry.generate_depth"]
output_path = self.storage_root / "region.png"
pixels = np.full((10, 10, 4), 255, dtype=np.uint8)
pixels.reshape(-1, 4)[:5, 3] = 25
Image.fromarray(pixels, mode="RGBA").save(output_path)
canvas, _selection, bounds = image_mask.create_selection_canvas(
source_path,
((-1, -1), (2, -1), (-1, 2)),
)
result = np.asarray(canvas)
self.assertEqual(depth.transparent_fraction(output_path), 0.05)
self.assertEqual(bounds, (-1, -1, 2, 2))
self.assertEqual(result.shape, (3, 3, 4))
self.assertEqual(tuple(result[2, 2]), (0, 0, 0, 0))
self.assertEqual(tuple(result[1, 1]), (20, 40, 60, 255))
def test_image_mask_crops_tightly_to_alpha(self):
import numpy as np
from PIL import Image
image_mask = sys.modules["server.selection.polygon_mask"]
pixels = np.zeros((4, 5, 4), dtype=np.uint8)
pixels[1:3, 2:4] = (20, 40, 60, 255)
result, bounds = image_mask.crop_to_alpha(
Image.fromarray(pixels, mode="RGBA"),
(10, 20, 15, 24),
)
result = np.asarray(result)
self.assertEqual(bounds, (12, 21, 14, 23))
self.assertEqual(result.shape, (2, 2, 4))
self.assertTrue(np.all(result[:, :, 3] == 255))
def test_mesh_refine_selection_keeps_source_size_and_selection_alpha(self):
cutout_job = importlib.import_module("server.selection.cutout_images")
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
source_path = self.storage_root / "source.png"
pixels = np.full((6, 8, 4), (20, 40, 60, 255), dtype=np.uint8)
Image.fromarray(pixels, mode="RGBA").save(source_path)
job_directory = self.storage_root / "job"
job_directory.mkdir()
context = SimpleNamespace(
directory=job_directory,
resource=lambda _name: SimpleNamespace(
directory=lambda _model: Path("model")
),
check_cancelled=lambda: None,
)
def passthrough(path, *_args, **_kwargs):
return Image.open(path).convert("RGBA"), 0.0, 0.0
output_path = job_directory / "image.0001.png"
with patch.object(
cutout_job,
"infer_foreground",
side_effect=passthrough,
):
cutout_job.refine_selection(
context,
{"refine_model": "BEN2_BASE", "device": "cpu"},
source_path,
((1, 1), (6, 1), (1, 5)),
output_path,
)
result = np.asarray(Image.open(output_path).convert("RGBA"))
self.assertEqual(result.shape, (6, 8, 4))
self.assertEqual(tuple(result[1, 1]), (20, 40, 60, 255))
self.assertEqual(tuple(result[5, 7]), (0, 0, 0, 0))
def test_mesh_cutout_images_crop_color_and_preserve_source_alpha(self):
cutout_job = importlib.import_module("server.selection.cutout_images")
import json
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
source_path = self.storage_root / "analysis-source.png"
selection_path = self.storage_root / "selection-source.png"
output_directory = self.storage_root / "cutout-images"
output_directory.mkdir()
source = np.full((6, 8, 4), (10, 80, 160, 255), dtype=np.uint8)
source[2, 3, 3] = 25
source[2, 4, 3] = 26
source[3, 5, 3] = 0
Image.fromarray(source, mode="RGBA").save(source_path)
selection = np.zeros((6, 8, 4), dtype=np.uint8)
selection[1:5, 2:7] = (200, 20, 40, 255)
selection[1, 2] = (200, 20, 40, 0)
Image.fromarray(selection, mode="RGBA").save(selection_path)
color_path, mask_path, analysis_path, bounds_path = (
cutout_job.write_images(
source_path,
selection_path,
output_directory,
)
)
color = np.asarray(Image.open(color_path).convert("RGBA"))
mask = np.asarray(Image.open(mask_path).convert("RGBA"))
analysis = np.asarray(Image.open(analysis_path).convert("RGBA"))
self.assertEqual(json.loads(bounds_path.read_text()), [2, 1, 7, 5])
self.assertEqual(color.shape, (4, 5, 4))
self.assertEqual(tuple(color[1, 1]), (10, 80, 160, 25))
self.assertEqual(int(mask[1, 1, 3]), 0)
self.assertEqual(tuple(analysis[1, 1]), (0, 0, 0, 0))
self.assertEqual(tuple(color[1, 2]), (10, 80, 160, 26))
self.assertGreater(int(mask[1, 2, 3]), 25)
self.assertEqual(int(mask[0, 0, 3]), 0)
self.assertEqual(tuple(color[0, 0]), (10, 80, 160, 255))
self.assertEqual(tuple(analysis[0, 0]), (0, 0, 0, 0))
self.assertEqual(tuple(analysis[1, 2]), (10, 80, 160, 26))
self.assertEqual(tuple(color[2, 3]), (10, 80, 160, 0))
self.assertEqual(int(mask[2, 3, 3]), 0)
self.assertEqual(tuple(analysis[2, 3]), (0, 0, 0, 0))
def test_mesh_unrefined_cutout_discards_selection_outside_the_source(self):
cutout_job = importlib.import_module("server.selection.cutout_images")
import json
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
source_path = self.storage_root / 'outside-source.png'
full_selection_path = self.storage_root / 'outside-selection-full.png'
output_directory = self.storage_root / 'outside-cutout'
output_directory.mkdir()
Image.new('RGBA', (8, 6), (10, 80, 160, 255)).save(source_path)
polygon = ((-2, 1), (10, 1), (10, 5), (-2, 5))
_selection_path, selection_bounds = (
cutout_job.write_selection(
source_path,
polygon,
full_selection_path,
)
)
color_path, mask_path, analysis_path, bounds_path = (
cutout_job.write_images(
source_path,
full_selection_path,
output_directory,
selection_bounds,
)
)
color = np.asarray(Image.open(color_path).convert('RGBA'))
mask = np.asarray(Image.open(mask_path).convert('RGBA'))
analysis = np.asarray(Image.open(analysis_path).convert('RGBA'))
self.assertEqual(selection_bounds, (-2, 1, 10, 5))
self.assertEqual(json.loads(bounds_path.read_text()), [0, 1, 8, 5])
self.assertEqual(color.shape, (4, 8, 4))
self.assertEqual(tuple(color[1, 0]), (10, 80, 160, 255))
self.assertEqual(tuple(color[1, 7]), (10, 80, 160, 255))
self.assertGreater(int(mask[1, 0, 3]), 127)
self.assertEqual(tuple(analysis[1, 0]), (10, 80, 160, 255))
def test_mesh_job_discards_unrefined_outside_selection_with_normal(self):
import json
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
source_path = self.storage_root / "mesh-source.png"
pixels = np.full((6, 8, 4), (20, 40, 60, 255), dtype=np.uint8)
Image.fromarray(pixels, mode="RGBA").save(source_path)
job_directory = self.storage_root / "mesh-job"
context = SimpleNamespace(
directory=job_directory,
progress=lambda *_args: None,
check_cancelled=lambda: None,
resource=lambda _name: object(),
)
normal = np.zeros((1, 4, 12, 3), dtype=np.float32)
normal[..., 0] = 1.0
prediction = SimpleNamespace(normal=normal)
received = []
received_images = []
def infer(_manager, parameters, images, **_kwargs):
received.append(parameters)
received_images.extend(images)
return prediction
with patch.object(mesh_job.moge2, "infer", side_effect=infer):
result = mesh_job.run(
context,
{
"input": str(source_path),
"shape": "BALLOON",
"polygon": ((-2, 1), (10, 1), (10, 5), (-2, 5)),
"mesh_detail": 64,
"invert": False,
"device": "cpu",
"model": "MOGE2_VITS_NORMAL",
"resolution_level": 5,
"normal_analysis": True,
},
)
color_path = job_directory / "image.0001.png"
color_pixels = np.asarray(Image.open(color_path).convert("RGBA"))
normal_pixels = np.asarray(
Image.open(job_directory / result["normal"]).convert("RGB")
)
self.assertEqual(
result,
{
"selection": "selection.png",
"bounds": "bounds.json",
"normal": "normal.png",
},
)
self.assertTrue(color_path.is_file())
self.assertFalse((job_directory / "contour.json").exists())
self.assertEqual(received[0]["model"], "MOGE2_VITS_NORMAL")
self.assertEqual(received_images, [job_directory / "analysis.png"])
analysis_pixels = np.asarray(
Image.open(received_images[0]).convert("RGBA")
)
self.assertEqual(
json.loads((job_directory / "bounds.json").read_text()),
[0, 1, 8, 5],
)
self.assertEqual(tuple(analysis_pixels[2, 2, :3]), (20, 40, 60))
self.assertGreater(int(analysis_pixels[2, 2, 3]), 240)
self.assertEqual(tuple(analysis_pixels[2, 0, :3]), (20, 40, 60))
self.assertEqual(tuple(color_pixels[2, 0, :3]), (20, 40, 60))
self.assertGreater(int(color_pixels[2, 0, 3]), 240)
self.assertEqual(tuple(color_pixels[2, 7, :3]), (20, 40, 60))
self.assertEqual(color_pixels.shape[:2], (4, 8))
self.assertEqual(analysis_pixels.shape[:2], (4, 8))
self.assertEqual(normal_pixels.shape[:2], (4, 8))
self.assertEqual(tuple(normal_pixels[0, 0]), (255, 127, 127))
self.assertEqual(tuple(normal_pixels[2, 0]), (255, 127, 127))
def test_cutout_limits_the_cropped_model_input_but_keeps_color_resolution(self):
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
source_path = self.storage_root / "limited-cutout-source.png"
Image.new("RGBA", (20, 10), (20, 40, 60, 255)).save(source_path)
job_directory = self.storage_root / "limited-cutout-job"
context = SimpleNamespace(
directory=job_directory,
progress=lambda *_args: None,
check_cancelled=lambda: None,
resource=lambda _name: object(),
)
prediction = SimpleNamespace(
normal=np.zeros((1, 2, 8, 3), dtype=np.float32)
)
inferred_sizes = []
def infer(_manager, _parameters, images, **_kwargs):
with Image.open(images[0]) as model_input:
inferred_sizes.append(model_input.size)
return prediction
with patch.object(mesh_job.moge2, "infer", side_effect=infer):
result = mesh_job.run(
context,
{
"input": str(source_path),
"shape": "BALLOON",
"polygon": ((2, 3), (18, 3), (18, 7), (2, 7)),
"mesh_detail": 64,
"max_input_size": 8,
"invert": False,
"device": "cpu",
"model": "MOGE2_VITS_NORMAL",
"resolution_level": 5,
"normal_analysis": True,
},
)
with Image.open(job_directory / "image.0001.png") as color:
color_size = color.size
with Image.open(job_directory / result["normal"]) as normal:
normal_size = normal.size
self.assertEqual(inferred_sizes, [(8, 2)])
self.assertEqual(color_size, (16, 4))
self.assertEqual(normal_size, (16, 4))
self.assertFalse((job_directory / "model-input.png").exists())
def test_mesh_job_skips_moge_normal_when_option_is_disabled(self):
from PIL import Image
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
source_path = self.storage_root / "mesh-no-normal-source.png"
Image.new("RGBA", (8, 6), (20, 40, 60, 255)).save(source_path)
job_directory = self.storage_root / "mesh-no-normal-job"
context = SimpleNamespace(
directory=job_directory,
progress=lambda *_args: None,
check_cancelled=lambda: None,
resource=lambda _name: object(),
)
with patch.object(mesh_job, "_write_camera_normal") as write_normal:
result = mesh_job.run(
context,
{
"input": str(source_path),
"shape": "BALLOON",
"polygon": ((0, 0), (8, 0), (8, 6), (0, 6)),
"mesh_detail": 64,
"normal_analysis": False,
"invert": False,
"device": "cpu",
},
)
self.assertEqual(
result,
{"selection": "selection.png", "bounds": "bounds.json"},
)
write_normal.assert_not_called()
self.assertFalse((job_directory / "normal.png").exists())
def test_mesh_depth_balloon_uses_metric_point_field(self):
from PIL import Image
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
source_path = self.storage_root / "mesh-depth-balloon-source.png"
Image.new("RGBA", (6, 6), (20, 40, 60, 255)).save(source_path)
job_directory = self.storage_root / "mesh-depth-balloon-job"
context = SimpleNamespace(
directory=job_directory,
progress=lambda *_args: None,
check_cancelled=lambda: None,
)
field_path = job_directory / "field.npz"
with patch.object(
mesh_job,
"_write_depth_field",
return_value=(field_path, None),
) as write_surface:
result = mesh_job.run(
context,
{
"input": str(source_path),
"shape": "DEPTH_BALLOON",
"polygon": ((1, 1), (5, 1), (5, 5), (1, 5)),
"normal_analysis": False,
"video_frames": 1,
"invert": False,
},
)
self.assertEqual(
result,
{
"selection": "selection.png",
"bounds": "bounds.json",
"field": "field.npz",
},
)
self.assertTrue(write_surface.called)
self.assertEqual(
write_surface.call_args.kwargs["selection_path"],
job_directory / "selection.png",
)
def test_moge_depth_and_normal_share_one_prediction(self):
def test_image_geometry_keeps_model_and_region_validity_separate(self):
import numpy as np
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
job_directory = self.storage_root / "shared-moge-prediction"
job_directory.mkdir()
color_path = job_directory / "analysis.png"
normal_path = job_directory / "normal.png"
prediction = object()
valid = np.ones((2, 3), dtype=bool)
depth = sys.modules["server.geometry.generate_depth"]
output_path = self.storage_root / "image_geometry.npz"
valid = np.asarray(((True, True, False), (True, True, True)))
calibration = np.asarray(((False, True, False), (False, False, True)))
frame = SimpleNamespace(
depth=np.ones((2, 3), dtype=np.float32),
depth=np.asarray(((1.0, 2.0, 3.0), (4.0, 5.0, 6.0))),
points=np.ones((2, 3, 3), dtype=np.float32),
intrinsics=np.eye(3, dtype=np.float32),
valid_mask=valid,
model_mask=valid,
)
context = SimpleNamespace(
directory=job_directory,
resource=lambda _name: object(),
check_cancelled=lambda: None,
)
parameters = {
"model": "MOGE2_VITS_NORMAL",
"device": "cpu",
"resolution_level": 5,
}
with (
patch.object(
mesh_job.moge2, "infer", return_value=prediction
) as infer,
patch.object(
mesh_job.moge2, "geometry_frame", return_value=frame
),
patch.object(mesh_job, "_write_field") as write_field,
patch.object(
mesh_job,
"_write_camera_normal_from_prediction",
return_value=normal_path,
) as write_normal,
):
field_path, written_normal_path = mesh_job._write_depth_field(
context,
parameters,
color_path,
include_normal=True,
)
infer.assert_called_once()
write_field.assert_called_once()
write_normal.assert_called_once_with(
prediction, color_path, normal_path
)
self.assertEqual(field_path, job_directory / "field.npz")
self.assertEqual(written_normal_path, normal_path)
def test_depth_calibration_mask_is_separate_from_model_validity(self):
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
selection_path = self.storage_root / "depth-selection.png"
selection = np.zeros((2, 3, 4), dtype=np.uint8)
selection[0, 1, 3] = 255
selection[1, 2, 3] = 255
Image.fromarray(selection, mode="RGBA").save(selection_path)
frame = SimpleNamespace(
valid_mask=np.asarray(
((True, True, True), (True, True, False)),
dtype=bool,
)
model_mask=None,
)
original_valid = frame.valid_mask.copy()
result = mesh_job._selection_calibration_mask(frame, selection_path)
depth._write_image_geometry_archive(frame, output_path, calibration, 0.05)
np.testing.assert_array_equal(
result,
((False, True, False), (False, False, False)),
)
np.testing.assert_array_equal(frame.valid_mask, original_valid)
def test_depth_field_writes_original_and_calibration_masks_separately(self):
import numpy as np
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
output_path = self.storage_root / "field.npz"
valid = np.asarray(((True, True), (True, False)), dtype=bool)
calibration = np.asarray(((False, True), (False, False)), dtype=bool)
depth = np.asarray(((2.0, 2.1), (2.2, 0.0)), dtype=np.float32)
points = np.zeros((2, 2, 3), dtype=np.float32)
points[..., 2] = depth
frame = SimpleNamespace(
depth=depth,
points=points,
intrinsics=np.eye(3, dtype=np.float32),
valid_mask=valid,
model_mask=valid,
)
mesh_job._write_field(frame, output_path, calibration)
with np.load(output_path, allow_pickle=False) as archive:
with np.load(output_path) as archive:
np.testing.assert_array_equal(archive["valid"], valid)
np.testing.assert_array_equal(
archive["calibration_valid"],
calibration,
np.testing.assert_array_equal(archive["calibration_valid"], calibration)
self.assertAlmostEqual(
float(archive["region_transparent_fraction"]),
0.05,
places=6,
)
def test_mesh_job_uses_upscale_source_only_for_analysis(self):
from PIL import Image
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
display_path = self.storage_root / "upscaled.png"
analysis_path = self.storage_root / "original.png"
Image.new("RGBA", (12, 8), (200, 20, 40, 255)).save(display_path)
Image.new("RGBA", (3, 2), (10, 80, 160, 255)).save(analysis_path)
job_directory = self.storage_root / "mesh-upscale-source-job"
def test_generate_region_data_only_returns_requested_outputs(self):
job = sys.modules["server.jobs.generate_region_data"]
input_path = self.storage_root / "input.png"
input_path.write_bytes(b"region")
context = SimpleNamespace(
directory=job_directory,
directory=self.storage_root / "job",
progress=lambda *_args: None,
check_cancelled=lambda: None,
)
field_path = job_directory / "field.npz"
archive = context.directory / "image_geometry.npz"
normal = context.directory / "normal.png"
with patch.object(
mesh_job,
"_write_depth_field",
return_value=(field_path, None),
) as write_surface:
result = mesh_job.run(
job,
"generate_depth",
return_value=(archive, normal),
) as generate:
result = job.run(
context,
{
"input": str(display_path),
"analysis_input": str(analysis_path),
"shape": "DEPTH_SURFACE",
"normal_analysis": False,
"video_frames": 1,
"input": str(input_path),
"output_depth": True,
"normal_mode": "OBJECT",
},
)
self.assertEqual(result, {"field": "field.npz"})
self.assertEqual(write_surface.call_args.args[2], analysis_path)
with Image.open(job_directory / "image.0001.png") as color:
self.assertEqual(color.size, (12, 8))
def test_mesh_masked_color_keeps_same_direction_overlap(self):
cutout_job = importlib.import_module("server.selection.cutout_images")
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
source_path = self.storage_root / "opaque-overlap.png"
output_path = self.storage_root / "cutout-overlap.png"
Image.new("RGBA", (12, 9), (20, 40, 60, 255)).save(source_path)
polygon = (
(1, 1),
(7, 1),
(7, 7),
(1, 7),
(1, 1),
(4, 1),
(10, 1),
(10, 7),
(4, 7),
(4, 1),
(1, 1),
self.assertEqual(
result,
{"image_geometry": "image_geometry.npz", "normal": "normal.png"},
)
generate.assert_called_once_with(
context,
{
"input": str(input_path),
"output_depth": True,
"normal_mode": "OBJECT",
},
input_path,
output_depth=True,
normal_mode="OBJECT",
)
cutout_job._write_masked_color(source_path, polygon, output_path)
alpha = np.asarray(Image.open(output_path).convert("RGBA"))[:, :, 3]
self.assertGreater(int(alpha[3, 2]), 200)
self.assertGreater(int(alpha[3, 5]), 200)
self.assertGreater(int(alpha[3, 8]), 200)
def test_mesh_masked_color_preserves_nonzero_winding_hole(self):
cutout_job = importlib.import_module("server.selection.cutout_images")
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
source_path = self.storage_root / "opaque-hole.png"
output_path = self.storage_root / "cutout-hole.png"
Image.new("RGBA", (10, 10), (20, 40, 60, 255)).save(source_path)
polygon = (
(1, 1),
(9, 1),
(9, 9),
(1, 9),
(1, 1),
(3, 3),
(3, 7),
(7, 7),
(7, 3),
(3, 3),
(1, 1),
def test_generate_region_data_refines_an_already_rasterized_region(self):
job = sys.modules["server.jobs.generate_region_data"]
input_path = self.storage_root / "input.png"
input_path.write_bytes(b"region")
directory = self.storage_root / "job"
context = SimpleNamespace(
directory=directory,
progress=lambda *_args: None,
check_cancelled=lambda: None,
)
cutout_job._write_masked_color(source_path, polygon, output_path)
alpha = np.asarray(Image.open(output_path).convert("RGBA"))[:, :, 3]
def generate(_context, _input, _model, _device, output):
output.write_bytes(b"refined")
return output, (2, 3, 8, 9)
self.assertGreater(int(alpha[2, 2]), 200)
self.assertLess(int(alpha[4, 4]), 32)
with patch.object(job, "generate_selection", side_effect=generate):
result = job.run(
context,
{
"input": str(input_path),
"refine_region": True,
"refine_model": "BEN2_BASE",
"device": "cpu",
},
)
def test_mesh_masked_color_antialiases_the_cutout_edge(self):
cutout_job = importlib.import_module("server.selection.cutout_images")
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.prepare_geometry_data"]
source_path = self.storage_root / "opaque.png"
output_path = self.storage_root / "cutout.png"
Image.new("RGBA", (16, 16), (20, 40, 60, 255)).save(source_path)
cutout_job._write_masked_color(
source_path,
((1, 1), (14, 1), (1, 14)),
output_path,
self.assertEqual(
result,
{"region": "region.png", "region_bounds": "region_bounds.json"},
)
self.assertEqual(
(directory / "region_bounds.json").read_text(encoding="utf-8"),
"[2,3,8,9]",
)
alpha = np.asarray(Image.open(output_path).convert("RGBA"))[:, :, 3]
self.assertTrue(np.any((alpha > 0) & (alpha < 255)))
if __name__ == "__main__":
unittest.main()

7
tests/test_worker.py
View File

@@ -927,6 +927,13 @@ class WorkerTest(unittest.TestCase):
self.assertEqual(result.mode, "RGBA")
self.assertEqual(result.size, (20, 12))
alpha = np.asarray(result.getchannel("A"))
expected_alpha = np.asarray(
Image.fromarray(pixels[:, :, 3], mode="L").resize(
result.size,
Image.Resampling.LANCZOS,
)
)
np.testing.assert_array_equal(alpha, expected_alpha)
self.assertLess(
int(alpha[:, 0].max()),
int(alpha[:, -1].min()),