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

Add new node asset scripts for geometry and shader operations

Browse Source 
- Implemented `build_mesh_plane.py` to create a geometry node group for mesh planes with customizable parameters.
- Created `build_shadeless.py` for a shader node group that displays color without scene lighting while preserving shadows and reflections.
- Added `validate.py` to validate the integrity of the newly created node groups and their inputs.
- Introduced `validate_cutout.py` to ensure the functionality of cutout node groups and their dependencies.
- Developed `validate_mesh_plane.py` to test the mesh plane creation and its modifier settings.
This commit is contained in:
马南
2026-08-30 00:17:29 +08:00
parent 7d39fe8c5e
commit 32bcd45d7b
168 changed files with 7213 additions and 13121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

613
tests/test_worker.py
View File

@@ -1,4 +1,3 @@
import importlib.util
import os
import sys
import tempfile
@@ -116,76 +115,29 @@ class WorkerTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.media = load_server_module("media.input")
cls.resolution = load_server_module("media.resolution")
cls.geometry = load_server_module("geometry")
cls.surface = load_server_module("geometry.surface")
cls.background = load_server_module("jobs.background")
cls.background = load_server_module("jobs.remove_background")
cls.ben2 = load_server_module("models.ben2")
cls.da3_model = load_server_module("models.da3")
cls.moge2_model = load_server_module("models.moge2")
cls.upscale_model = load_server_module("models.upscale")
cls.onnx_upscale = load_server_module("models.onnx_upscale")
cls.surface_outputs = load_server_module("outputs.surface")
cls.debug_da3_outputs = load_server_module("outputs.debug_da3")
cls.debug_moge2_outputs = load_server_module("outputs.debug_moge2")
cls.common_job = load_server_module("jobs.common")
cls.depth_job = load_server_module("jobs.depth")
cls.surface_job = load_server_module("jobs.surface")
cls.debug_da3_job = load_server_module("jobs.debug_da3")
cls.debug_moge2_job = load_server_module("jobs.debug_moge2")
cls.debug_ben2_job = load_server_module("jobs.debug_ben2")
cls.debug_upscale_job = load_server_module(
"jobs.debug_upscale"
)
cls.debug_jobs = load_server_module("jobs.debug_predictions")
cls.upscale_job = load_server_module("jobs.upscale")
cls.frames_job = load_server_module("jobs.frames")
cls.foreground_job = load_server_module("jobs.foreground")
cls.download = load_server_module("model_download")
def test_inference_progress_omits_single_frame_count(self):
def test_upscale_progress_omits_single_frame_count(self):
self.assertEqual(
self.common_job.frame_progress_message(
"Estimating MoGe-2 geometry", 0, 1
),
"Estimating MoGe-2 geometry",
self.upscale_job.frame_progress_message("Upscaling image", 0, 1),
"Upscaling image",
)
self.assertEqual(
self.common_job.frame_progress_message(
"Estimating DA3 geometry", 0, 1
),
"Estimating DA3 geometry",
self.upscale_job.frame_progress_message("Upscaling image", 1, 3),
"Upscaling image (2/3)",
)
self.assertEqual(
self.common_job.frame_progress_message(
"Estimating DA3 geometry", 1, 3
),
"Estimating DA3 geometry (2/3)",
)
self.assertEqual(
self.common_job.frame_progress_message(
"Estimating MoGe-2 depth", 0, 1
),
"Estimating MoGe-2 depth",
)
self.assertEqual(
self.common_job.frame_progress_message(
"Estimating DA3 depth", 0, 1
),
"Estimating DA3 depth",
)
def test_surface_has_one_model_independent_builder(self):
self.assertTrue(callable(self.surface.build_surface))
self.assertTrue(callable(self.surface_outputs.write_surface_archive))
self.assertFalse(hasattr(self.surface_outputs, "write_surface_texture"))
self.assertFalse(hasattr(self.surface, "export_da3_surface"))
self.assertFalse(hasattr(self.surface, "export_moge2_surface"))
def test_model_adapters_only_normalize_surface_input(self):
self.assertTrue(callable(self.da3_model.surface_frame))
self.assertTrue(callable(self.moge2_model.surface_frame))
self.assertFalse(hasattr(self.da3_model, "generate_surface"))
self.assertFalse(hasattr(self.moge2_model, "generate_surface"))
def test_collect_images_is_sorted_and_filters_suffixes(self):
with tempfile.TemporaryDirectory() as directory:
@@ -740,35 +692,6 @@ class WorkerTest(unittest.TestCase):
for _image, options in inference_options:
self.assertEqual(options, {})
def test_extract_frames_copies_selected_frames(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
input_directory = root / "frames"
input_directory.mkdir()
for index in range(1, 6):
(input_directory / f"frame-{index:04d}.jpg").touch()
context = FakeJobContext(root / "result", "extract-frames")
self.frames_job.run(
context,
{
"input": str(input_directory),
"video_frames": 3,
"start_frame": 2,
},
)
result_directory = context.directory
self.assertEqual(
sorted(
path.name
for path in result_directory.iterdir()
if path.is_file()
),
["image.0001.jpg", "image.0002.jpg", "image.0003.jpg"],
)
self.assertEqual(context.status["progress"], 1.0)
self.assertIn("3 frame", context.status["message"])
def test_da3_debug_saves_only_direct_api_prediction(self):
prediction = SimpleNamespace(
depth=np.full((1, 2, 2), 2.0, dtype=np.float32),
@@ -802,7 +725,7 @@ class WorkerTest(unittest.TestCase):
},
):
context = FakeJobContext(root / "result", "debug-da3")
self.debug_da3_job.run(
self.debug_jobs.run_da3(
context,
{
"input": str(input_path),
@@ -848,11 +771,11 @@ class WorkerTest(unittest.TestCase):
(frames / "b.png").touch()
context = FakeJobContext(root / "result", "debug-moge2")
with patch.object(
self.debug_moge2_job,
"infer_one",
self.debug_jobs,
"infer_moge2",
return_value=prediction,
) as infer_one:
self.debug_moge2_job.run(
self.debug_jobs.run_moge2(
context,
{
"input": str(frames),
@@ -880,11 +803,11 @@ class WorkerTest(unittest.TestCase):
input_path.touch()
context = FakeJobContext(root / "result", "debug-ben2")
with patch.object(
self.debug_ben2_job,
self.debug_jobs,
"infer_foreground",
return_value=(Foreground(), 1.0, 2.0),
) as infer_foreground:
self.debug_ben2_job.run(
self.debug_jobs.run_ben2(
context,
{
"input": str(input_path),
@@ -910,11 +833,11 @@ class WorkerTest(unittest.TestCase):
(frames / "b.png").touch()
context = FakeJobContext(root / "result", "debug-upscale")
with patch.object(
self.debug_upscale_job,
"infer_one",
self.debug_jobs,
"infer_upscale",
return_value=Image.new("RGB", (8, 8), (10, 20, 30)),
) as infer_one:
self.debug_upscale_job.run(
self.debug_jobs.run_upscale(
context,
{
"input": str(frames),
@@ -967,7 +890,7 @@ class WorkerTest(unittest.TestCase):
(
"INFO",
"Upscale model: REALESRGAN_GENERAL_WDN_X4V3; "
"maximum output size: 4096 px",
"maximum AI input size: 2048 px",
),
context.logs,
)
@@ -1009,24 +932,19 @@ class WorkerTest(unittest.TestCase):
int(alpha[:, -1].min()),
)
def test_upscale_input_is_limited_before_four_x_inference(self):
def test_upscale_input_uses_the_global_long_edge_limit(self):
from PIL import Image
source = Image.new("RGB", (2000, 1000))
limited = self.upscale_model._limited_input(source, 4096)
source = Image.new("RGB", (3000, 1500))
limited = self.resolution.limit_image(source, 2048)
self.assertEqual(limited.size, (1024, 512))
square = self.upscale_model._limited_input(
Image.new("RGB", (3000, 3000)),
8192,
)
self.assertLessEqual(
square.width * square.height * 16,
self.upscale_model.MAX_OUTPUT_PIXELS,
self.assertEqual(limited.size, (2048, 1024))
self.assertEqual(
self.resolution.limit_image(Image.new("RGB", (800, 600)), 2048).size,
(800, 600),
)
with self.assertRaisesRegex(ValueError, "between 1 and 8192"):
self.upscale_model._limited_input(source, 9000)
self.resolution.limit_image(source, 9000)
def test_upscale_resource_error_releases_session_and_is_readable(self):
from PIL import Image
@@ -1057,495 +975,20 @@ class WorkerTest(unittest.TestCase):
):
with self.assertRaisesRegex(
RuntimeError,
"Reduce Max Output Size or select CPU",
"Reduce Maximum AI Input Size or select CPU",
):
self.upscale_model.infer_one(
models,
{
"model": "REALESRGAN_GENERAL_WDN_X4V3",
"device": "directml",
"max_output_size": 4096,
"max_input_size": 2048,
},
path,
)
self.assertTrue(models.released)
def test_da3_removes_background_after_depth_using_original_input(self):
order = []
inference_inputs = []
fake_da3 = ModuleType("server.models.onnx_da3")
def fake_infer(_session, images, **_options):
order.append("depth")
inference_inputs.append(images)
return SimpleNamespace(
depth=np.full((1, 2, 2), 2.0, dtype=np.float32),
conf=np.ones((1, 2, 2), dtype=np.float32),
extrinsics=np.eye(4, dtype=np.float32)[None, :3],
intrinsics=np.eye(3, dtype=np.float32)[None],
processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8),
)
fake_da3.infer = fake_infer
fake_runtime = ModuleType("server.models.onnx_runtime")
fake_runtime.provider_label = lambda _session: "CPU"
class FakeForeground:
def convert(self, _mode):
return self
def save(self, path, **_options):
Path(path).touch()
def fake_foreground(image, *_args, **_kwargs):
order.append("background")
self.assertEqual(image, input_path)
return FakeForeground(), 1.0, 2.0
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
input_path = root / "input.png"
input_path.touch()
with patch.dict(
sys.modules,
{
"server.models.onnx_da3": fake_da3,
"server.models.onnx_runtime": fake_runtime,
},
), patch.object(
self.foreground_job,
"infer_foreground",
side_effect=fake_foreground,
):
context = FakeJobContext(
root / "result",
"generate-depth",
)
self.depth_job.run(
context,
{
"input": str(input_path),
"model": "DA3",
"remove_background_model": "BEN2",
"video_frames": 1,
"export_format": "mini_npz",
"device": "cpu",
"process_res": 518,
"process_res_method": "upper_bound_resize",
"ref_view_strategy": "first",
"inference_mode": "per_frame",
},
)
self.assertEqual(order, ["depth", "background"])
self.assertEqual(inference_inputs, [[str(input_path)]])
def test_reconstruct_copies_frames_for_multi_frame_input(self):
inference_inputs = []
fake_da3 = ModuleType("server.models.onnx_da3")
fake_da3.create_session = lambda _model_dir, _device: object()
def fake_infer(_session, images, **_options):
inference_inputs.append(images)
return SimpleNamespace(
depth=np.full((1, 2, 2), 2.0, dtype=np.float32),
conf=np.ones((1, 2, 2), dtype=np.float32),
extrinsics=np.eye(4, dtype=np.float32)[None, :3],
intrinsics=np.eye(3, dtype=np.float32)[None],
processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8),
)
fake_da3.infer = fake_infer
fake_runtime = ModuleType("server.models.onnx_runtime")
fake_runtime.provider_label = lambda _session: "CPU"
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
input_directory = root / "frames"
input_directory.mkdir()
(input_directory / "frame-0001.jpg").touch()
(input_directory / "frame-0002.jpg").touch()
with patch.dict(
sys.modules,
{
"server.models.onnx_da3": fake_da3,
"server.models.onnx_runtime": fake_runtime,
},
):
context = FakeJobContext(
root / "result",
"generate-depth",
)
self.depth_job.run(
context,
{
"input": str(input_directory),
"model": "DA3",
"video_frames": 16,
"export_format": "mini_npz",
"device": "cpu",
"process_res": 518,
"process_res_method": "upper_bound_resize",
"ref_view_strategy": "first",
"conf_thresh_percentile": 40.0,
"num_max_points": 1000,
"inference_mode": "per_frame",
},
)
result_directory = context.directory
self.assertTrue(
(result_directory / "image.0001.jpg").is_file()
)
self.assertTrue(
(result_directory / "image.0002.jpg").is_file()
)
self.assertEqual(len(inference_inputs), 2)
self.assertTrue(all(len(paths) == 1 for paths in inference_inputs))
with np.load(result_directory / "results.npz") as result:
self.assertEqual(result["depth"].shape, (2, 2, 2))
self.assertTrue(
(result_directory / "results.npz").is_file()
)
self.assertFalse(
(result_directory / "exports").exists()
)
self.assertEqual(context.status["progress"], 1.0)
def test_mesh_confidence_uses_percentile(self):
depth = np.ones((1, 2, 2), dtype=np.float32)
confidence = np.array([[[1.0, 2.0], [3.0, 4.0]]])
threshold = self.geometry.confidence_percentile_threshold(
depth,
confidence,
50.0,
)
self.assertEqual(threshold, 2.5)
def test_zero_mesh_confidence_disables_confidence_filter(self):
depth = np.ones((2, 2), dtype=np.float32)
confidence = np.full((2, 2), np.nan, dtype=np.float32)
x, y = np.meshgrid([0.0, 1.0], [0.0, 1.0])
points = np.stack((x, y, depth), axis=-1)
threshold = self.geometry.confidence_percentile_threshold(
depth,
confidence,
0.0,
)
faces = self.geometry.mesh_triangle_faces(
points,
depth,
confidence,
threshold,
filter_edges=False,
max_edge_length=0.0,
)
self.assertIsNone(threshold)
self.assertEqual(len(faces), 2)
def test_mesh_filters_triangles_independently(self):
depth = np.ones((2, 2), dtype=np.float32)
confidence = np.ones((2, 2), dtype=np.float32)
confidence[1, 0] = 0.0
x, y = np.meshgrid(
np.arange(2, dtype=np.float32),
np.arange(2, dtype=np.float32),
)
points = np.stack((x, y, depth), axis=-1)
faces = self.geometry.mesh_triangle_faces(
points,
depth,
confidence,
confidence_threshold=0.5,
filter_edges=False,
max_edge_length=0.0,
)
np.testing.assert_array_equal(faces, [[0, 3, 1]])
def test_mesh_background_mask_combines_with_confidence(self):
depth = np.ones((2, 3), dtype=np.float32)
confidence = np.ones((2, 3), dtype=np.float32)
x, y = np.meshgrid(
np.arange(3, dtype=np.float32),
np.arange(2, dtype=np.float32),
)
points = np.stack((x, y, depth), axis=-1)
foreground = np.ones((2, 3), dtype=bool)
foreground[:, 0] = False
faces = self.geometry.mesh_triangle_faces(
points,
depth,
confidence,
confidence_threshold=None,
filter_edges=False,
max_edge_length=0.0,
vertex_mask=foreground,
)
np.testing.assert_array_equal(faces, [[1, 4, 5], [1, 5, 2]])
def test_depth_edge_filter_ignores_smooth_slope(self):
smooth = np.tile(
np.linspace(2.0, 2.04, 5, dtype=np.float32),
(3, 1),
)
discontinuous = smooth.copy()
discontinuous[:, 3:] += 1.0
self.assertFalse(self.geometry.depth_edge_mask(smooth).any())
self.assertTrue(self.geometry.depth_edge_mask(discontinuous)[:, 2:4].all())
def test_higher_edge_culling_allows_less_stretch(self):
limits = [
self.geometry.edge_culling_stretch_limit(value)
for value in (0.0, 50.0, 100.0)
]
self.assertEqual(limits, [10000.0, 2.0, 1.0])
def test_stretched_edge_culling_uses_camera_space_edge_ratio(self):
flat_points = np.asarray(
[
[[0.0, 0.0, 1.0], [1.0, 0.0, 1.0]],
[[0.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
],
dtype=np.float32,
)
stretched_points = flat_points.copy()
stretched_points[0, 1] = (2.0, 0.0, 2.0)
triangles = np.asarray([[0, 2, 3], [0, 3, 1]], dtype=np.int32)
flat_mask = self.geometry.triangle_edge_stretch_mask(
flat_points,
triangles,
1.0,
)
loose_mask = self.geometry.triangle_edge_stretch_mask(
stretched_points,
triangles,
2.0,
)
strict_mask = self.geometry.triangle_edge_stretch_mask(
stretched_points,
triangles,
1.0,
)
self.assertTrue(flat_mask.all())
self.assertTrue(loose_mask.all())
self.assertFalse(strict_mask.all())
def test_mesh_faces_apply_confidence_threshold_and_face_direction(self):
points = np.array(
[
[[0.0, 0.0, 1.0], [1.0, 0.0, 1.0], [2.0, 0.0, 1.0]],
[[0.0, 1.0, 1.0], [1.0, 1.0, 1.0], [2.0, 1.0, 1.0]],
],
dtype=np.float32,
)
depth = np.ones((2, 3), dtype=np.float32)
confidence = np.ones((2, 3), dtype=np.float32)
confidence[:, 0] = 0.0
faces = self.geometry.mesh_triangle_faces(
points,
depth,
confidence,
confidence_threshold=0.5,
filter_edges=False,
max_edge_length=0.0,
)
np.testing.assert_array_equal(faces, [[1, 4, 5], [1, 5, 2]])
vertices = points.reshape(-1, 3)[faces[0]]
normal = np.cross(
vertices[1] - vertices[0],
vertices[2] - vertices[0],
)
self.assertLess(normal[2], 0.0)
def test_mesh_percentile_reduces_face_count(self):
depth = np.ones((3, 3), dtype=np.float32)
confidence = np.arange(1, 10, dtype=np.float32).reshape(3, 3)
x, y = np.meshgrid(
np.arange(3, dtype=np.float32),
np.arange(3, dtype=np.float32),
)
points = np.stack((x, y, depth), axis=-1)
face_counts = []
for percentile in (0.0, 50.0, 100.0):
threshold = self.geometry.confidence_percentile_threshold(
depth,
confidence,
percentile,
)
faces = self.geometry.mesh_triangle_faces(
points,
depth,
confidence,
confidence_threshold=threshold,
filter_edges=False,
max_edge_length=0.0,
)
face_counts.append(len(faces))
self.assertEqual(face_counts, [8, 2, 0])
def test_mesh_compaction_removes_unreferenced_vertices(self):
vertices = np.arange(18, dtype=np.float32).reshape(6, 3)
faces = np.asarray([[1, 3, 5]], dtype=np.int32)
compact_vertices, compact_faces = self.geometry.compact_mesh(
vertices,
faces,
)
np.testing.assert_array_equal(
compact_vertices,
vertices[[1, 3, 5]],
)
np.testing.assert_array_equal(compact_faces, [[0, 1, 2]])
def test_surface_background_requires_every_vertex_to_be_foreground(self):
faces = np.asarray(
((0, 1, 4), (0, 4, 3), (1, 2, 5), (1, 5, 4)),
dtype=np.int32,
)
foreground = np.asarray(
((True, True, False), (False, True, False)),
dtype=bool,
)
filtered = self.geometry.remove_background_faces(faces, foreground)
np.testing.assert_array_equal(filtered, faces[[0]])
def test_surface_background_mask_uses_half_alpha_threshold(self):
from PIL import Image
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
for alpha, expected in (
(0, False),
(127, False),
(128, True),
(255, True),
):
path = root / f"alpha-{alpha}.png"
Image.new("RGBA", (2, 2), (10, 20, 30, alpha)).save(path)
mask = self.geometry.foreground_mesh_mask(path, (2, 2))
self.assertEqual(mask.dtype, np.bool_)
self.assertTrue(np.all(mask == expected))
def test_image_mesh_data_keeps_camera_geometry_and_uv(self):
depth = np.full((3, 3), 2.0, dtype=np.float32)
confidence = np.ones((3, 3), dtype=np.float32)
intrinsics = np.asarray(
[[2.0, 0.0, 1.0], [0.0, 2.0, 1.0], [0.0, 0.0, 1.0]],
dtype=np.float32,
)
frame = self.surface.GeometryFrame(
depth=depth,
points=self.geometry.unproject_camera_depth(depth, intrinsics),
intrinsics=intrinsics,
valid_mask=np.ones(depth.shape, dtype=bool),
confidence=confidence,
)
mesh = self.surface.build_surface(frame)
with tempfile.TemporaryDirectory() as directory:
result_path = self.surface_outputs.write_surface_archive(
mesh,
Path(directory) / "mesh.npz",
)
with np.load(result_path) as result:
self.assertEqual(result["vertices"].shape, (9, 3))
self.assertEqual(result["faces"].shape, (8, 3))
self.assertEqual(result["uv"].shape, (9, 2))
self.assertEqual(float(result["reference_depth"]), 2.0)
self.assertNotIn("postprocess_method", result.files)
self.assertEqual(result["point_depth"].shape, (9,))
self.assertEqual(
result["face_confidence"].shape,
(8,),
)
self.assertEqual(
set(result.files),
{
"vertices",
"faces",
"uv",
"intrinsics",
"image_size",
"reference_depth",
"point_depth",
"face_confidence",
"face_depth_discontinuity",
},
)
np.testing.assert_allclose(result["uv"][0], [0.0, 1.0])
np.testing.assert_allclose(result["uv"][-1], [1.0, 0.0])
self.assertNotIn("normal_map", result.files)
def test_moge_surface_uses_the_same_archive_without_normal_data(self):
depth = np.full((3, 3), 2.0, dtype=np.float32)
intrinsics = np.asarray(
[[2.0, 0.0, 1.0], [0.0, 2.0, 1.0], [0.0, 0.0, 1.0]],
dtype=np.float32,
)
frame = self.surface.GeometryFrame(
depth=depth,
points=self.geometry.unproject_camera_depth(depth, intrinsics),
intrinsics=intrinsics,
valid_mask=np.ones(depth.shape, dtype=bool),
model_mask=np.asarray(
(
(False, False, False),
(True, True, True),
(True, True, True),
),
dtype=bool,
),
)
mesh = self.surface.build_surface(frame)
with tempfile.TemporaryDirectory() as directory:
path = self.surface_outputs.write_surface_archive(
mesh,
Path(directory) / "mesh.npz",
)
with np.load(path) as result:
self.assertIn("face_mask", result.files)
self.assertNotIn("face_confidence", result.files)
self.assertNotIn("normal", result.files)
self.assertNotIn("normal_map", result.files)
def test_moge_surface_normal_is_exported_as_a_texture(self):
from PIL import Image
normal = np.asarray(
[[[0.0, 0.0, -1.0]]],
dtype=np.float32,
)
with tempfile.TemporaryDirectory() as directory:
path = self.surface_outputs.write_normal_texture(
normal,
Path(directory) / "normal.png",
)
with Image.open(path) as image:
self.assertEqual(image.mode, "RGB")
self.assertEqual(image.size, (1, 1))
red, green, blue = image.getpixel((0, 0))
self.assertIn(red, {127, 128})
self.assertIn(green, {127, 128})
self.assertEqual(blue, 255)
def test_da3_debug_writer_keeps_api_values_without_derived_attributes(self):
prediction = SimpleNamespace(
depth=np.ones((1, 2, 3), dtype=np.float32),