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

Add comprehensive tests for cutout and mesh functionalities

Browse Source 
- Introduced `test_cutout.py` to validate cutout pie menu behavior, image lasso options, and mesh detail calculations.
- Added `test_mesh.py` to ensure correct tangent normal rasterization, plane source topology, and depth surface evaluations.
- Created `test_mesh_blender.py` to verify material node groups, tangent baking, and local cutout building from Blender images.
- Implemented tests for handling alpha regions and ensuring geometry nodes maintain expected behavior.
This commit is contained in:
马南
2026-08-26 10:42:29 +08:00
parent 0cf04a13c6
commit 2b79a342b0
46 changed files with 5554 additions and 645 deletions
Download Patch File Download Diff File Expand all files Collapse all files

380
tests/test_server.py
View File

@@ -8,7 +8,6 @@ from pathlib import Path
from types import ModuleType, SimpleNamespace
from unittest.mock import patch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
@@ -42,6 +41,8 @@ class ServerTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.server_module = load_server_module()
importlib.import_module("server.jobs.balloon")
importlib.import_module("server.jobs.mesh")
from blendjob import JobServer
from server.model_manager import ModelManager, ModelSessionCache
@@ -50,6 +51,7 @@ class ServerTest(unittest.TestCase):
cls.ModelManager, cls.ModelSessionCache = ModelManager, ModelSessionCache
cls.server = cls.server_module.server
cls.app_module = sys.modules["server.app"]
def setUp(self):
self.jobs_temp = tempfile.TemporaryDirectory()
self.addCleanup(self.jobs_temp.cleanup)
@@ -89,8 +91,7 @@ class ServerTest(unittest.TestCase):
{
"loaded": {
"ben2": "",
"da3": "",
"moge2": "",
"geometry": "",
"upscale": "",
}
},
@@ -99,9 +100,7 @@ class ServerTest(unittest.TestCase):
self.assertNotIn("model_catalog", sys.modules)
def create_cache(self):
cache = self.ModelSessionCache.__new__(
self.ModelSessionCache
)
cache = self.ModelSessionCache.__new__(self.ModelSessionCache)
cache.ben2_model = None
cache.ben2_key = None
cache.da3_model = None
@@ -114,9 +113,7 @@ class ServerTest(unittest.TestCase):
def create_server(self):
source = self.server_module.server
server = self.JobServer(
"AnyImage Job Server", storage_root=self.storage_root
)
server = self.JobServer("AnyImage Job Server", storage_root=self.storage_root)
server.handlers.update(source.handlers)
server.resource_factories.update(source.resource_factories)
server.bind(self.storage_root)
@@ -168,14 +165,19 @@ class ServerTest(unittest.TestCase):
self.assertIs(first, second)
self.assertEqual(loads, [("model", "cuda")])
self.assertEqual(cache.snapshot()["da3"], "model")
self.assertEqual(cache.snapshot()["geometry"], "model")
def test_model_cache_keeps_only_the_most_recent_family(self):
def test_model_cache_keeps_three_categories(self):
cache = self.create_cache()
cache.ben2_model = object()
ben2 = object()
old_moge = object()
upscale = object()
cache.ben2_model = ben2
cache.ben2_key = ("ben2", "cuda")
cache.moge_model = object()
cache.moge_key = ("moge", "cuda")
cache.moge_model = old_moge
cache.moge_key = ("old-moge", "cuda")
cache.upscale_model = upscale
cache.upscale_key = ("upscale", "cuda")
device = SimpleNamespace(type="cuda")
da3 = object()
@@ -187,18 +189,21 @@ class ServerTest(unittest.TestCase):
with patch.object(cache, "_load_da3", side_effect=load_da3):
self.assertIs(cache.get_da3("da3", device), da3)
self.assertIs(cache.ben2_model, ben2)
self.assertIsNone(cache.moge_model)
self.assertIsNone(cache.ben2_model)
moge = object()
self.assertIs(cache.upscale_model, upscale)
self.assertEqual(cache.snapshot()["geometry"], "da3")
new_moge = object()
fake_moge = ModuleType("server.models.onnx_moge2")
fake_moge.create_session = lambda _model_dir, _device: moge
fake_moge.create_session = lambda _model_dir, _device: new_moge
with patch.dict(sys.modules, {"server.models.onnx_moge2": fake_moge}):
self.assertIs(cache.get_moge2("moge", "cuda"), moge)
self.assertIs(cache.get_moge2("new-moge", "cuda"), new_moge)
self.assertIs(cache.ben2_model, ben2)
self.assertIsNone(cache.da3_model)
self.assertIsNone(cache.ben2_model)
self.assertIs(cache.upscale_model, upscale)
self.assertEqual(cache.snapshot()["geometry"], "new-moge")
def test_upscale_model_is_reused_for_matching_runtime_key(self):
loads = []
@@ -254,13 +259,11 @@ class ServerTest(unittest.TestCase):
{
"download-model",
"download-required-models",
"analyze-balloon",
"refine-image-mask",
"remove-background",
"upscale-image",
"extract-frames",
"generate-depth",
"generate-surface",
"generate-mesh",
"debug-da3",
"debug-moge2",
"debug-ben2",
@@ -438,6 +441,54 @@ class ServerTest(unittest.TestCase):
self.assertEqual(tuple(result[0, 0]), (20, 40, 60, 64))
self.assertEqual(tuple(result[3, 3]), (0, 0, 0, 0))
def test_polygon_selection_mask_uses_even_odd_fill_for_self_intersection(self):
image_mask = sys.modules["server.jobs.image_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.jobs.image_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):
import numpy as np
from PIL import Image
@@ -475,6 +526,288 @@ class ServerTest(unittest.TestCase):
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):
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.mesh"]
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):
return Image.open(path).convert("RGBA"), 0.0, 0.0
output_path = job_directory / "image.0001.png"
with patch.object(mesh_job, "remove_image_background", side_effect=passthrough):
mesh_job._write_refined_color(
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_job_keeps_original_ai_input_separate_from_geometry_mask(self):
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.mesh"]
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, 6, 8, 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), (7, 1), (7, 5), (2, 5)),
"mesh_detail": 64,
"invert": False,
"device": "cpu",
"normal_model": "MOGE2_VITS_NORMAL",
"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, {"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, [source_path])
self.assertEqual(tuple(color_pixels[3, 4, :3]), (20, 40, 60))
self.assertGreater(int(color_pixels[3, 4, 3]), 240)
self.assertLess(int(color_pixels[0, 0, 3]), 5)
self.assertEqual(tuple(normal_pixels[0, 0]), (255, 127, 127))
self.assertEqual(tuple(normal_pixels[3, 4]), (255, 127, 127))
def test_mesh_job_skips_moge_normal_when_option_is_disabled(self):
from PIL import Image
mesh_job = sys.modules["server.jobs.mesh"]
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, {})
write_normal.assert_not_called()
self.assertFalse((job_directory / "normal.png").exists())
def test_mesh_depth_balloon_analyzes_full_source_then_applies_selection(self):
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.mesh"]
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"
job_directory.mkdir()
context = SimpleNamespace(
directory=job_directory,
check_cancelled=lambda: None,
resource=lambda _name: object(),
)
received = []
def infer(_manager, _parameters, images, **_kwargs):
received.extend(images)
depth = np.linspace(1.0, 2.0, 36, dtype=np.float32).reshape(6, 6)
normal = np.zeros((1, 6, 6, 3), dtype=np.float32)
normal[..., 0] = 1.0
return SimpleNamespace(
depth=depth[None],
mask=np.ones((1, 6, 6), dtype=bool),
normal=normal,
)
with patch.object(mesh_job.moge2, "infer", side_effect=infer):
depth_path, normal_path = mesh_job._write_balloon_depth(
context,
{"device": "cpu"},
source_path,
((2, 2), (10, 2), (2, 10)),
selection_size=(12, 12),
include_normal=True,
)
self.assertEqual(received, [source_path])
with np.load(depth_path, allow_pickle=False) as data:
self.assertEqual(tuple(data["origin"]), (0, 0))
self.assertEqual(tuple(data["crop_size"]), (6, 6))
self.assertFalse(bool(data["valid"][0, 0]))
self.assertTrue(bool(data["valid"][2, 2]))
normal = np.asarray(Image.open(normal_path).convert("RGB"))
self.assertEqual(tuple(normal[0, 0]), (255, 127, 127))
self.assertEqual(tuple(normal[2, 2]), (255, 127, 127))
def test_mesh_job_uses_upscale_source_only_for_analysis(self):
from PIL import Image
mesh_job = sys.modules["server.jobs.mesh"]
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"
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_surface_depth",
return_value=(field_path, None),
) as write_surface:
result = mesh_job.run(
context,
{
"input": str(display_path),
"analysis_input": str(analysis_path),
"shape": "DEPTH_SURFACE",
"normal_analysis": False,
"video_frames": 1,
},
)
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):
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.mesh"]
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),
)
mesh_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):
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.mesh"]
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),
)
mesh_job._write_masked_color(source_path, polygon, output_path)
alpha = np.asarray(Image.open(output_path).convert("RGBA"))[:, :, 3]
self.assertGreater(int(alpha[2, 2]), 200)
self.assertLess(int(alpha[4, 4]), 32)
def test_mesh_masked_color_antialiases_the_cutout_edge(self):
import numpy as np
from PIL import Image
mesh_job = sys.modules["server.jobs.mesh"]
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)
mesh_job._write_masked_color(
source_path,
((1, 1), (14, 1), (1, 14)),
output_path,
)
alpha = np.asarray(Image.open(output_path).convert("RGBA"))[:, :, 3]
def test_balloon_job_sends_the_masked_crop_to_ben2(self):
import numpy as np
from PIL import Image
@@ -587,6 +920,7 @@ class ServerTest(unittest.TestCase):
np.testing.assert_allclose(farther[3], original[3], atol=2e-4)
np.testing.assert_allclose(rescaled[0], original[0], atol=2e-4)
np.testing.assert_allclose(rescaled[3], original[3], atol=2e-4)
def test_balloon_depth_job_sends_masked_crop_only_to_moge(self):
import numpy as np
from PIL import Image