nan
bdbde1a10d
- Renamed "Base Plane Depth" to "Reference Depth" across validate_cutout.py for consistency. - Introduced a new specification for Cutout mesh sampling to ensure source pixel spacing is preserved and limits are applied correctly. - Unified naming conventions for image editing results to maintain consistency across Remove Background, Upscale, and Crop operations. - Updated the handling of Refine Selection to ensure it retains the original bounds and does not apply unnecessary cropping. - Enhanced material generation logic to connect image color directly to IOR Level, improving the rendering of dark areas. - Established clear requirements for geometry depth controls and image result materials to ensure consistent behavior across different components. - Removed outdated logic and ensured all changes are documented and tested for regression.
880 lines
31 KiB
Python
880 lines
31 KiB
Python
import importlib.util
|
|
import json
|
|
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 ServerRuntimeTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.server_module = load_server_module()
|
|
importlib.import_module("server.geometry.prediction_artifacts")
|
|
importlib.import_module("server.jobs.cutout")
|
|
from blendjob import JobServer
|
|
from server.model_manager import ModelManager, ModelSessionCache
|
|
from server.models import onnx_runtime
|
|
|
|
cls.JobServer = JobServer
|
|
|
|
cls.ModelManager, cls.ModelSessionCache = ModelManager, ModelSessionCache
|
|
cls.onnx_runtime = onnx_runtime
|
|
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_inference_cleanup_keeps_every_cached_session(self):
|
|
class Session:
|
|
def get_providers(self):
|
|
return ["CUDAExecutionProvider"]
|
|
|
|
def run(self, _names, _feeds, run_options=None):
|
|
self.run_options = run_options
|
|
return [object()]
|
|
|
|
class RunOptions:
|
|
def add_run_config_entry(self, _key, _value):
|
|
return None
|
|
|
|
runtime = ModuleType("onnxruntime")
|
|
runtime.RunOptions = RunOptions
|
|
sessions = {
|
|
"ben2": Session(),
|
|
"da3": Session(),
|
|
"moge2": Session(),
|
|
"upscale": Session(),
|
|
}
|
|
ben2 = ModuleType("server.models.onnx_ben2")
|
|
ben2.create_session = lambda _directory, _device: sessions["ben2"]
|
|
da3 = ModuleType("server.models.onnx_da3")
|
|
da3.create_session = lambda _directory, _device: sessions["da3"]
|
|
moge2 = ModuleType("server.models.onnx_moge2")
|
|
moge2.create_session = lambda _directory, _device: sessions["moge2"]
|
|
upscale = ModuleType("server.models.onnx_upscale")
|
|
upscale.create_session = lambda _directory, _device: sessions["upscale"]
|
|
cache = self.create_cache()
|
|
|
|
def assert_reused(first, getter, release_name):
|
|
with (
|
|
patch.dict(sys.modules, {"onnxruntime": runtime}),
|
|
patch.object(
|
|
cache,
|
|
release_name,
|
|
wraps=getattr(cache, release_name),
|
|
) as release,
|
|
):
|
|
self.onnx_runtime.run_session(first, ["output"], {"input": 1})
|
|
second = getter()
|
|
self.assertIs(first, second)
|
|
release.assert_not_called()
|
|
|
|
with patch.dict(
|
|
sys.modules,
|
|
{
|
|
"server.models.onnx_ben2": ben2,
|
|
"server.models.onnx_da3": da3,
|
|
"server.models.onnx_moge2": moge2,
|
|
"server.models.onnx_upscale": upscale,
|
|
},
|
|
):
|
|
first_ben2, _load_ms = cache.get_ben2("ben2", "cuda")
|
|
assert_reused(
|
|
first_ben2,
|
|
lambda: cache.get_ben2("ben2", "cuda")[0],
|
|
"release_ben2",
|
|
)
|
|
|
|
first_da3 = cache.get_da3("da3", "cuda")
|
|
assert_reused(
|
|
first_da3,
|
|
lambda: cache.get_da3("da3", "cuda"),
|
|
"release_geometry",
|
|
)
|
|
|
|
first_moge2 = cache.get_moge2("moge2", "cuda")
|
|
assert_reused(
|
|
first_moge2,
|
|
lambda: cache.get_moge2("moge2", "cuda"),
|
|
"release_geometry",
|
|
)
|
|
|
|
first_upscale = cache.get_upscale("upscale", "cuda")
|
|
assert_reused(
|
|
first_upscale,
|
|
lambda: cache.get_upscale("upscale", "cuda"),
|
|
"release_upscale",
|
|
)
|
|
|
|
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-selection",
|
|
"remove-background",
|
|
"upscale-image",
|
|
"generate-depth-plane-geometry",
|
|
"generate-cutout-artifacts",
|
|
"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_depth_artifacts_store_texture_and_compact_metadata(self):
|
|
import numpy as np
|
|
|
|
depth = importlib.import_module("server.geometry.depth_texture")
|
|
GeometryFrame = importlib.import_module(
|
|
"server.models.geometry"
|
|
).GeometryFrame
|
|
texture_path = self.storage_root / "z-depth.exr"
|
|
metadata_path = self.storage_root / "depth.json"
|
|
validity = np.asarray(
|
|
((1.0, 1.0, 0.0), (1.0, 1.0, 1.0)),
|
|
dtype=np.float32,
|
|
)
|
|
calibration = np.asarray(((False, True, True), (False, False, True)))
|
|
frame = GeometryFrame(
|
|
depth=np.asarray(
|
|
((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)),
|
|
dtype=np.float32,
|
|
),
|
|
validity=validity,
|
|
points=np.ones((2, 3, 3), dtype=np.float32),
|
|
intrinsics=np.eye(3, dtype=np.float32),
|
|
)
|
|
|
|
depth.write_z_depth_texture(frame, texture_path)
|
|
depth.write_depth_metadata(frame, calibration, metadata_path)
|
|
|
|
self.assertTrue(texture_path.is_file())
|
|
self.assertEqual(texture_path.read_bytes()[:4], b"v/1\x01")
|
|
metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
|
|
self.assertEqual(metadata["image_size"], [3, 2])
|
|
self.assertEqual(metadata["intrinsics"], np.eye(3).tolist())
|
|
self.assertAlmostEqual(metadata["reference_depth"], 4.0)
|
|
|
|
def test_object_normal_texture_matches_depth_surface_axes(self):
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
normal_texture = importlib.import_module("server.geometry.normal_texture")
|
|
GeometryFrame = importlib.import_module(
|
|
"server.models.geometry"
|
|
).GeometryFrame
|
|
output_path = self.storage_root / "object-normal.png"
|
|
camera_normals = np.asarray(
|
|
[[
|
|
(0.0, 0.0, -1.0),
|
|
(1.0, 0.0, 0.0),
|
|
(0.0, 1.0, 0.0),
|
|
(0.0, 0.0, -1.0),
|
|
]],
|
|
dtype=np.float32,
|
|
)
|
|
valid = np.asarray([[True, True, True, False]], dtype=bool)
|
|
frame = GeometryFrame(
|
|
depth=np.ones((1, 4), dtype=np.float32),
|
|
validity=np.ones((1, 4), dtype=np.float32),
|
|
intrinsics=np.eye(3, dtype=np.float32),
|
|
normal=camera_normals,
|
|
)
|
|
|
|
normal_texture.write_camera_normal_texture(
|
|
frame,
|
|
valid,
|
|
(0, 0),
|
|
(4, 1),
|
|
(4, 1),
|
|
output_path,
|
|
normal_mode="OBJECT",
|
|
)
|
|
|
|
with Image.open(output_path) as opened:
|
|
encoded = np.asarray(opened, dtype=np.uint8)
|
|
np.testing.assert_array_equal(
|
|
encoded,
|
|
np.asarray(
|
|
[[
|
|
(0, 127, 127),
|
|
(127, 0, 127),
|
|
(127, 127, 0),
|
|
(0, 128, 128),
|
|
]],
|
|
dtype=np.uint8,
|
|
),
|
|
)
|
|
|
|
def test_moge_model_input_composites_rgba_over_black(self):
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
artifacts = sys.modules["server.geometry.prediction_artifacts"]
|
|
first_path = self.storage_root / "first.png"
|
|
second_path = self.storage_root / "second.png"
|
|
first = np.asarray(
|
|
[[(240, 120, 60, 0), (200, 100, 50, 128), (9, 8, 7, 255)]],
|
|
dtype=np.uint8,
|
|
)
|
|
second = first.copy()
|
|
second[0, 0, :3] = (1, 2, 3)
|
|
Image.fromarray(first, mode="RGBA").save(first_path)
|
|
Image.fromarray(second, mode="RGBA").save(second_path)
|
|
|
|
first_output = artifacts._limited_model_input(
|
|
first_path,
|
|
self.storage_root / "first-model.png",
|
|
2048,
|
|
)
|
|
second_output = artifacts._limited_model_input(
|
|
second_path,
|
|
self.storage_root / "second-model.png",
|
|
2048,
|
|
)
|
|
|
|
with Image.open(first_output) as opened:
|
|
first_rgb = np.asarray(opened, dtype=np.uint8)
|
|
with Image.open(second_output) as opened:
|
|
second_rgb = np.asarray(opened, dtype=np.uint8)
|
|
np.testing.assert_array_equal(first_rgb, second_rgb)
|
|
np.testing.assert_array_equal(
|
|
first_rgb,
|
|
np.asarray([[(0, 0, 0), (100, 50, 25), (9, 8, 7)]], dtype=np.uint8),
|
|
)
|
|
|
|
def test_geometry_writers_reject_missing_or_non_finite_fields(self):
|
|
import numpy as np
|
|
|
|
depth = importlib.import_module("server.geometry.depth_texture")
|
|
normal = importlib.import_module("server.geometry.normal_texture")
|
|
GeometryFrame = importlib.import_module(
|
|
"server.models.geometry"
|
|
).GeometryFrame
|
|
base = {
|
|
"depth": np.ones((1, 2), dtype=np.float32),
|
|
"validity": np.ones((1, 2), dtype=np.float32),
|
|
"intrinsics": np.eye(3, dtype=np.float32),
|
|
}
|
|
with self.assertRaisesRegex(ValueError, "point field"):
|
|
depth.write_vector_depth_texture(
|
|
GeometryFrame(**base),
|
|
self.storage_root / "missing-points.exr",
|
|
)
|
|
point_frame = GeometryFrame(
|
|
**base,
|
|
points=np.asarray([[(0.0, 0.0, 1.0), (np.nan, 0.0, 1.0)]], dtype=np.float32),
|
|
)
|
|
with self.assertRaisesRegex(ValueError, "Point field"):
|
|
depth.write_vector_depth_texture(
|
|
point_frame,
|
|
self.storage_root / "nan-points.exr",
|
|
)
|
|
depth_frame = GeometryFrame(
|
|
**{**base, "depth": np.asarray(((1.0, np.inf),), dtype=np.float32)}
|
|
)
|
|
with self.assertRaisesRegex(ValueError, "Depth field"):
|
|
depth.write_z_depth_texture(
|
|
depth_frame,
|
|
self.storage_root / "inf-depth.exr",
|
|
)
|
|
with self.assertRaisesRegex(ValueError, "normal field"):
|
|
normal.write_camera_normal_texture(
|
|
GeometryFrame(**base),
|
|
np.ones((1, 2), dtype=bool),
|
|
(0, 0),
|
|
(2, 1),
|
|
(2, 1),
|
|
self.storage_root / "missing-normal.png",
|
|
)
|
|
normal_frame = GeometryFrame(
|
|
**base,
|
|
normal=np.asarray(
|
|
[[(0.0, 0.0, 1.0), (0.0, np.nan, 1.0)]],
|
|
dtype=np.float32,
|
|
),
|
|
)
|
|
with self.assertRaisesRegex(ValueError, "Normal field"):
|
|
normal.write_camera_normal_texture(
|
|
normal_frame,
|
|
np.ones((1, 2), dtype=bool),
|
|
(0, 0),
|
|
(2, 1),
|
|
(2, 1),
|
|
self.storage_root / "nan-normal.png",
|
|
)
|
|
self.assertFalse(any(self.storage_root.glob("*.exr")))
|
|
self.assertFalse(any(self.storage_root.glob("*-normal.png")))
|
|
|
|
def test_moge_artifact_orchestration_uses_one_geometry_frame(self):
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
artifacts = sys.modules["server.geometry.prediction_artifacts"]
|
|
GeometryFrame = importlib.import_module(
|
|
"server.models.geometry"
|
|
).GeometryFrame
|
|
input_path = self.storage_root / "selection.png"
|
|
Image.new("RGBA", (2, 2), (255, 128, 64, 255)).save(input_path)
|
|
output_directory = self.storage_root / "job"
|
|
output_directory.mkdir()
|
|
context = SimpleNamespace(
|
|
directory=output_directory,
|
|
resource=lambda _name: object(),
|
|
progress=lambda *_args: None,
|
|
check_cancelled=lambda: None,
|
|
)
|
|
frame = GeometryFrame(
|
|
depth=np.ones((2, 2), dtype=np.float32),
|
|
validity=np.ones((2, 2), dtype=np.float32),
|
|
intrinsics=np.eye(3, dtype=np.float32),
|
|
normal=np.zeros((2, 2, 3), dtype=np.float32),
|
|
points=np.ones((2, 2, 3), dtype=np.float32),
|
|
)
|
|
|
|
with patch.object(artifacts.moge2, "infer", return_value=[frame]) as infer:
|
|
result = artifacts.generate_moge2_artifacts(
|
|
context,
|
|
{"max_input_size": 2048},
|
|
input_path,
|
|
depth_mode="VECTOR",
|
|
normal_mode="OBJECT",
|
|
)
|
|
|
|
infer.assert_called_once()
|
|
self.assertEqual(
|
|
set(result),
|
|
{"vector_depth", "depth_metadata", "object_normal"},
|
|
)
|
|
self.assertFalse((output_directory / "model-input.png").exists())
|
|
self.assertTrue(all(path.is_file() for path in result.values()))
|
|
self.assertTrue(infer.call_args.kwargs["include_points"])
|
|
|
|
def test_refine_selection_preserves_the_full_alpha_canvas(self):
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
refine = importlib.import_module("server.selection.refine_selection")
|
|
input_path = self.storage_root / "input.png"
|
|
output_path = self.storage_root / "output.png"
|
|
source = np.zeros((3, 4, 4), dtype=np.uint8)
|
|
source[..., :3] = (32, 64, 96)
|
|
source[..., 3] = np.asarray(
|
|
((0, 64, 128, 255), (255, 128, 64, 0), (0, 0, 255, 255)),
|
|
dtype=np.uint8,
|
|
)
|
|
Image.fromarray(source, mode="RGBA").save(input_path)
|
|
foreground = Image.fromarray(source, mode="RGBA")
|
|
context = SimpleNamespace(
|
|
resource=lambda _name: SimpleNamespace(directory=lambda _model: "model"),
|
|
check_cancelled=lambda: None,
|
|
)
|
|
|
|
with patch.object(
|
|
refine,
|
|
"infer_foreground",
|
|
return_value=(foreground, 0.0, 0.0),
|
|
) as infer:
|
|
result = refine.refine_selection(
|
|
context,
|
|
input_path,
|
|
"BEN2_BASE",
|
|
"cpu",
|
|
output_path,
|
|
)
|
|
|
|
self.assertEqual(result, output_path)
|
|
infer.assert_called_once()
|
|
self.assertTrue(infer.call_args.kwargs["preserve_input_alpha"])
|
|
with Image.open(output_path) as output:
|
|
self.assertEqual(output.size, (4, 3))
|
|
np.testing.assert_array_equal(np.asarray(output.convert("RGBA")), source)
|
|
|
|
def test_generate_cutout_artifacts_only_returns_requested_outputs(self):
|
|
job = sys.modules["server.jobs.cutout"]
|
|
input_path = self.storage_root / "input.png"
|
|
input_path.write_bytes(b"selection")
|
|
context = SimpleNamespace(
|
|
directory=self.storage_root / "job",
|
|
progress=lambda *_args: None,
|
|
check_cancelled=lambda: None,
|
|
)
|
|
vector_depth = context.directory / "vector-depth.exr"
|
|
metadata = context.directory / "depth.json"
|
|
normal = context.directory / "object-normal.png"
|
|
with patch.object(
|
|
job,
|
|
"generate_moge2_artifacts",
|
|
return_value={
|
|
"vector_depth": vector_depth,
|
|
"depth_metadata": metadata,
|
|
"object_normal": normal,
|
|
},
|
|
) as generate:
|
|
result = job.run(
|
|
context,
|
|
{
|
|
"input": str(input_path),
|
|
"depth_mode": "VECTOR",
|
|
"normal_mode": "OBJECT",
|
|
},
|
|
)
|
|
|
|
self.assertEqual(
|
|
result,
|
|
{
|
|
"vector_depth": "vector-depth.exr",
|
|
"depth_metadata": "depth.json",
|
|
"object_normal": "object-normal.png",
|
|
},
|
|
)
|
|
generate.assert_called_once_with(
|
|
context,
|
|
{
|
|
"input": str(input_path),
|
|
"depth_mode": "VECTOR",
|
|
"normal_mode": "OBJECT",
|
|
},
|
|
input_path,
|
|
depth_mode="VECTOR",
|
|
normal_mode="OBJECT",
|
|
)
|
|
|
|
def test_generate_cutout_artifacts_refines_the_full_bounds_image(self):
|
|
job = sys.modules["server.jobs.cutout"]
|
|
input_path = self.storage_root / "input.png"
|
|
input_path.write_bytes(b"selection")
|
|
directory = self.storage_root / "job"
|
|
context = SimpleNamespace(
|
|
directory=directory,
|
|
progress=lambda *_args: None,
|
|
check_cancelled=lambda: None,
|
|
)
|
|
|
|
def generate(_context, _input, _model, _device, output):
|
|
output.write_bytes(b"refined")
|
|
return output
|
|
|
|
with patch.object(job, "refine_selection", side_effect=generate):
|
|
result = job.run(
|
|
context,
|
|
{
|
|
"input": str(input_path),
|
|
"refine_selection": True,
|
|
"refine_model": "BEN2_BASE",
|
|
"device": "cpu",
|
|
},
|
|
)
|
|
|
|
self.assertEqual(
|
|
result,
|
|
{"selection_image": "selection.png"},
|
|
)
|
|
self.assertFalse((directory / "selection_bounds.json").exists())
|
|
|
|
def test_refined_cutout_artifacts_share_the_full_refine_image(self):
|
|
job = sys.modules["server.jobs.cutout"]
|
|
input_path = self.storage_root / "input.png"
|
|
input_path.write_bytes(b"selection")
|
|
directory = self.storage_root / "job"
|
|
context = SimpleNamespace(
|
|
directory=directory,
|
|
progress=lambda *_args: None,
|
|
check_cancelled=lambda: None,
|
|
)
|
|
refined_path = directory / "selection.png"
|
|
|
|
def refine(_context, _input, _model, _device, output):
|
|
output.write_bytes(b"refined")
|
|
return output
|
|
|
|
with (
|
|
patch.object(job, "refine_selection", side_effect=refine),
|
|
patch.object(
|
|
job,
|
|
"generate_moge2_artifacts",
|
|
return_value={
|
|
"z_depth": directory / "z-depth.exr",
|
|
"depth_metadata": directory / "depth.json",
|
|
"tangent_normal": directory / "tangent-normal.png",
|
|
},
|
|
) as generate,
|
|
):
|
|
result = job.run(
|
|
context,
|
|
{
|
|
"input": str(input_path),
|
|
"refine_selection": True,
|
|
"refine_model": "BEN2_BASE",
|
|
"device": "cpu",
|
|
"depth_mode": "Z",
|
|
"normal_mode": "TANGENT",
|
|
},
|
|
)
|
|
|
|
self.assertEqual(result["selection_image"], "selection.png")
|
|
generate.assert_called_once_with(
|
|
context,
|
|
{
|
|
"input": str(input_path),
|
|
"refine_selection": True,
|
|
"refine_model": "BEN2_BASE",
|
|
"device": "cpu",
|
|
"depth_mode": "Z",
|
|
"normal_mode": "TANGENT",
|
|
},
|
|
refined_path,
|
|
depth_mode="Z",
|
|
normal_mode="TANGENT",
|
|
)
|