nan
32bcd45d7b
- 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.
1037 lines
38 KiB
Python
1037 lines
38 KiB
Python
import importlib.util
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import ModuleType, SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def load_server_module():
|
|
server_directory = PROJECT_ROOT / "src" / "anyimage" / "server"
|
|
specification = importlib.util.spec_from_file_location(
|
|
"server",
|
|
server_directory / "__init__.py",
|
|
submodule_search_locations=[str(server_directory)],
|
|
)
|
|
module = importlib.util.module_from_spec(specification)
|
|
sys.modules["server"] = module
|
|
specification.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class FakeCuda:
|
|
class OutOfMemoryError(RuntimeError):
|
|
pass
|
|
|
|
@staticmethod
|
|
def is_available():
|
|
return False
|
|
|
|
@staticmethod
|
|
def empty_cache():
|
|
return None
|
|
|
|
|
|
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")
|
|
from blendjob import JobServer
|
|
from server.model_manager import ModelManager, ModelSessionCache
|
|
|
|
cls.JobServer = JobServer
|
|
|
|
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)
|
|
self.storage_root = Path(self.jobs_temp.name)
|
|
|
|
def test_required_model_download_skips_ready_model(self):
|
|
manager = self.ModelManager(Path("."))
|
|
downloads = []
|
|
context = SimpleNamespace(progress=lambda *_args: None)
|
|
with (
|
|
patch.object(
|
|
manager,
|
|
"ready",
|
|
side_effect=lambda key: key == "BEN2_BASE",
|
|
),
|
|
patch.object(
|
|
manager,
|
|
"_download",
|
|
side_effect=lambda _context, key, _progress: downloads.append(key),
|
|
),
|
|
):
|
|
result = manager.download_missing(
|
|
context,
|
|
("MOGE2_VITS_NORMAL", "BEN2_BASE"),
|
|
)
|
|
|
|
self.assertEqual(downloads, ["MOGE2_VITS_NORMAL"])
|
|
self.assertEqual(result, {"downloaded": ["MOGE2_VITS_NORMAL"]})
|
|
|
|
def test_models_are_owned_by_the_server_resource(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
manager = self.ModelManager(Path(directory))
|
|
snapshot = manager.snapshot()
|
|
self.assertEqual(
|
|
snapshot,
|
|
{
|
|
"loaded": {
|
|
"ben2": "",
|
|
"geometry": "",
|
|
"upscale": "",
|
|
}
|
|
},
|
|
)
|
|
self.assertIn("server.model_catalog", sys.modules)
|
|
self.assertNotIn("model_catalog", sys.modules)
|
|
|
|
def create_cache(self):
|
|
cache = self.ModelSessionCache.__new__(self.ModelSessionCache)
|
|
cache.ben2_model = None
|
|
cache.ben2_key = None
|
|
cache.da3_model = None
|
|
cache.da3_key = None
|
|
cache.moge_model = None
|
|
cache.moge_key = None
|
|
cache.upscale_model = None
|
|
cache.upscale_key = None
|
|
return cache
|
|
|
|
def create_server(self):
|
|
source = self.server_module.server
|
|
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)
|
|
return server
|
|
|
|
def wait_for_job(self, context, timeout=2.0):
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if context._snapshot()["state"] in {"succeeded", "failed", "cancelled"}:
|
|
return context._snapshot()
|
|
time.sleep(0.01)
|
|
self.fail(f"Job did not finish: {context._snapshot()}")
|
|
|
|
def test_ben2_model_is_reused_for_matching_runtime_key(self):
|
|
loads = []
|
|
|
|
ben2 = ModuleType("server.models.onnx_ben2")
|
|
ben2.create_session = lambda model_dir, device: (
|
|
loads.append((model_dir, device)) or object()
|
|
)
|
|
cache = self.create_cache()
|
|
device = SimpleNamespace(type="cuda")
|
|
with patch.dict(sys.modules, {"server.models.onnx_ben2": ben2}):
|
|
first, _load_ms = cache.get_ben2("model", device)
|
|
second, cached_load_ms = cache.get_ben2("model", device)
|
|
|
|
self.assertIs(first, second)
|
|
self.assertEqual(loads, [("model", "cuda")])
|
|
self.assertEqual(cache.snapshot()["ben2"], "model")
|
|
self.assertEqual(cached_load_ms, 0.0)
|
|
|
|
def test_da3_model_is_reused_for_matching_runtime_key(self):
|
|
loads = []
|
|
|
|
da3 = ModuleType("server.models.onnx_da3")
|
|
da3.create_session = lambda model_dir, device: (
|
|
loads.append((model_dir, device)) or object()
|
|
)
|
|
cache = self.create_cache()
|
|
device = SimpleNamespace(type="cuda")
|
|
with patch.dict(
|
|
sys.modules,
|
|
{
|
|
"server.models.onnx_da3": da3,
|
|
},
|
|
):
|
|
first = cache.get_da3("model", device)
|
|
second = cache.get_da3("model", device)
|
|
|
|
self.assertIs(first, second)
|
|
self.assertEqual(loads, [("model", "cuda")])
|
|
self.assertEqual(cache.snapshot()["geometry"], "model")
|
|
|
|
def test_model_cache_keeps_three_categories(self):
|
|
cache = self.create_cache()
|
|
ben2 = object()
|
|
old_moge = object()
|
|
upscale = object()
|
|
cache.ben2_model = ben2
|
|
cache.ben2_key = ("ben2", "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()
|
|
|
|
def load_da3(_model_dir, _device, key):
|
|
cache.da3_model = da3
|
|
cache.da3_key = key
|
|
return da3
|
|
|
|
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.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: new_moge
|
|
with patch.dict(sys.modules, {"server.models.onnx_moge2": fake_moge}):
|
|
self.assertIs(cache.get_moge2("new-moge", "cuda"), new_moge)
|
|
|
|
self.assertIs(cache.ben2_model, ben2)
|
|
self.assertIsNone(cache.da3_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 = []
|
|
upscale = ModuleType("server.models.onnx_upscale")
|
|
upscale.create_session = lambda model_dir, device: (
|
|
loads.append((model_dir, device)) or object()
|
|
)
|
|
cache = self.create_cache()
|
|
|
|
with patch.dict(
|
|
sys.modules,
|
|
{"server.models.onnx_upscale": upscale},
|
|
):
|
|
first = cache.get_upscale("model", "cpu")
|
|
second = cache.get_upscale("model", "cpu")
|
|
|
|
self.assertIs(first, second)
|
|
self.assertEqual(loads, [("model", "cpu")])
|
|
self.assertEqual(cache.snapshot()["upscale"], "model")
|
|
|
|
def test_dedicated_server_runs_only_one_job_at_a_time(self):
|
|
server = self.JobServer("Test", storage_root=self.storage_root)
|
|
first_started = threading.Event()
|
|
release_first = threading.Event()
|
|
order = []
|
|
|
|
@server.job("example")
|
|
def example(_context, parameters):
|
|
order.append(parameters["name"])
|
|
if parameters["name"] == "first":
|
|
first_started.set()
|
|
release_first.wait(1.0)
|
|
|
|
first = server.submit("example", {"name": "first"}, job_id="job")
|
|
self.assertTrue(first_started.wait(1.0))
|
|
second = server.submit("example", {"name": "second"}, job_id="other")
|
|
|
|
self.assertIsNotNone(first)
|
|
self.assertIsNotNone(second)
|
|
self.assertTrue(server.snapshot("instance")["busy"])
|
|
self.assertEqual(server.snapshot("instance")["queued_jobs"], 1)
|
|
self.assertNotIn("clients", server.snapshot("instance"))
|
|
release_first.set()
|
|
self.wait_for_job(first)
|
|
self.wait_for_job(second)
|
|
self.assertEqual(order, ["first", "second"])
|
|
server.close()
|
|
|
|
def test_server_registers_tasks_instead_of_model_names(self):
|
|
server = self.create_server()
|
|
self.assertEqual(
|
|
set(server.handlers),
|
|
{
|
|
"download-model",
|
|
"download-required-models",
|
|
"refine-image-mask",
|
|
"remove-background",
|
|
"upscale-image",
|
|
"prepare-geometry-data",
|
|
"debug-da3",
|
|
"debug-moge2",
|
|
"debug-ben2",
|
|
"debug-upscale",
|
|
},
|
|
)
|
|
|
|
def test_queued_job_can_be_cancelled_before_it_runs(self):
|
|
server = self.JobServer("Test", storage_root=self.storage_root)
|
|
first_started = threading.Event()
|
|
release_first = threading.Event()
|
|
executed = []
|
|
|
|
@server.job("example")
|
|
def example(_context, parameters):
|
|
executed.append(parameters["name"])
|
|
if parameters["name"] == "first":
|
|
first_started.set()
|
|
release_first.wait(1.0)
|
|
|
|
first = server.submit("example", {"name": "first"})
|
|
self.assertTrue(first_started.wait(1.0))
|
|
second = server.submit("example", {"name": "second"})
|
|
|
|
self.assertIs(server.cancel(second.job_id), second)
|
|
release_first.set()
|
|
self.wait_for_job(first)
|
|
second_status = self.wait_for_job(second)
|
|
|
|
self.assertEqual(second_status["state"], "cancelled")
|
|
self.assertEqual(executed, ["first"])
|
|
server.close()
|
|
|
|
def test_cancelling_job_keeps_server_and_model_cache_alive(self):
|
|
fake_models = SimpleNamespace(
|
|
snapshot=lambda: {"loaded": {"ben2": "BEN2_BASE", "da3": ""}},
|
|
clear=lambda: None,
|
|
close=lambda: None,
|
|
)
|
|
with patch.object(self.app_module, "ModelManager", return_value=fake_models):
|
|
server = self.create_server()
|
|
started = threading.Event()
|
|
|
|
@server.job("cancellable")
|
|
def cancellable(context, _parameters):
|
|
started.set()
|
|
while not context._is_cancelled():
|
|
time.sleep(0.001)
|
|
context.check_cancelled()
|
|
|
|
context = server.submit("cancellable", {}, job_id="job")
|
|
self.assertTrue(started.wait(1.0))
|
|
self.assertIsNotNone(server.cancel("job"))
|
|
self.assertTrue(context._is_cancelled())
|
|
status = self.wait_for_job(context)
|
|
|
|
self.assertEqual(status["state"], "cancelled")
|
|
self.assertFalse(server.shutdown_event.is_set())
|
|
self.assertEqual(
|
|
server.snapshot("instance")["resources"]["model_manager"]["loaded"]["ben2"],
|
|
"BEN2_BASE",
|
|
)
|
|
|
|
def test_job_progress_is_kept_in_server_memory(self):
|
|
server = self.JobServer("Test", storage_root=self.storage_root)
|
|
|
|
@server.job("example")
|
|
def example(context, _parameters):
|
|
context.progress(0.5, "Half way")
|
|
|
|
context = server.submit("example", {})
|
|
self.wait_for_job(context)
|
|
|
|
self.assertRegex(context.job_id, r"^[0-9a-f]{32}$")
|
|
self.assertEqual(context.directory.name, context.job_id)
|
|
self.assertEqual(context._snapshot()["state"], "succeeded")
|
|
self.assertEqual(context._snapshot()["progress"], 1.0)
|
|
|
|
def test_runtime_storage_root_owns_the_jobs_directory(self):
|
|
storage_root = self.storage_root
|
|
server = self.JobServer("Test")
|
|
server.bind(storage_root)
|
|
|
|
@server.job("example")
|
|
def example(_context, _parameters):
|
|
return {"value": 42}
|
|
|
|
context = server.submit("example", {})
|
|
self.wait_for_job(context)
|
|
|
|
self.assertEqual(server.jobs_directory, storage_root / "jobs")
|
|
self.assertEqual(context.storage_root, storage_root.resolve())
|
|
self.assertEqual(context._snapshot()["result"], {"value": 42})
|
|
|
|
def test_job_server_owns_and_clears_shared_resources(self):
|
|
events = []
|
|
resource = SimpleNamespace(
|
|
snapshot=lambda: {"loaded": "model"},
|
|
clear=lambda: events.append("clear"),
|
|
close=lambda: events.append("close"),
|
|
)
|
|
server = self.JobServer("Test", storage_root=self.storage_root)
|
|
server.add_resource("model_manager", resource)
|
|
|
|
@server.job("example")
|
|
def example(context, _parameters):
|
|
self.assertIs(context.resource("model_manager"), resource)
|
|
release.wait(1.0)
|
|
|
|
release = threading.Event()
|
|
context = server.submit("example", {}, job_id="job")
|
|
self.assertFalse(server.clear_resource("model_manager"))
|
|
release.set()
|
|
self.wait_for_job(context)
|
|
self.assertTrue(server.clear_resource("model_manager"))
|
|
server.close()
|
|
|
|
self.assertEqual(
|
|
server.snapshot("instance")["resources"]["model_manager"]["loaded"],
|
|
"model",
|
|
)
|
|
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):
|
|
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)
|
|
|
|
canvas, _selection, bounds = image_mask.create_selection_canvas(
|
|
source_path,
|
|
((-1, -1), (2, -1), (-1, 2)),
|
|
)
|
|
result = np.asarray(canvas)
|
|
|
|
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):
|
|
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)
|
|
frame = SimpleNamespace(
|
|
depth=np.ones((2, 3), dtype=np.float32),
|
|
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,
|
|
)
|
|
)
|
|
|
|
original_valid = frame.valid_mask.copy()
|
|
result = mesh_job._selection_calibration_mask(frame, selection_path)
|
|
|
|
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:
|
|
np.testing.assert_array_equal(archive["valid"], valid)
|
|
np.testing.assert_array_equal(
|
|
archive["calibration_valid"],
|
|
calibration,
|
|
)
|
|
|
|
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"
|
|
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(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):
|
|
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),
|
|
)
|
|
|
|
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),
|
|
)
|
|
|
|
cutout_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):
|
|
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,
|
|
)
|
|
alpha = np.asarray(Image.open(output_path).convert("RGBA"))[:, :, 3]
|
|
|
|
self.assertTrue(np.any((alpha > 0) & (alpha < 255)))
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|