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.jobs.balloon") importlib.import_module("server.jobs.mesh") 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): self.assertFalse((PROJECT_ROOT / "src" / "anyimage" / "models.py").exists()) 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", "extract-frames", "generate-mesh", "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_server_dispatches_frame_extraction_without_models(self): fake_models = SimpleNamespace( snapshot=lambda: {"loaded": {"ben2": "", "da3": ""}}, clear=lambda: None, close=lambda: None, ) with patch.object(self.app_module, "ModelManager", return_value=fake_models): server = self.create_server() 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() context = server.submit( "extract-frames", { "input": str(input_directory), "video_frames": 16, "start_frame": 1, }, job_id="job", ) status = self.wait_for_job(context) self.assertTrue((context.directory / "image.0001.jpg").is_file()) self.assertEqual(status["state"], "succeeded") def test_balloon_lasso_input_crops_and_clears_outside_rgba(self): import numpy as np from PIL import Image balloon = sys.modules["server.jobs.balloon"] with tempfile.TemporaryDirectory() as directory: root = Path(directory) source_path = root / "source.png" output_path = root / "lasso.png" pixels = np.full((6, 6, 4), (20, 40, 60, 255), dtype=np.uint8) pixels[1, 1, 3] = 64 Image.fromarray(pixels, mode="RGBA").save(source_path) _path, _selection, origin, source_size = balloon.create_lasso_input( source_path, ((1, 1), (5, 1), (1, 5)), output_path, ) result = np.asarray(Image.open(output_path).convert("RGBA")) self.assertEqual(origin, (1, 1)) self.assertEqual(source_size, (6, 6)) self.assertEqual(result.shape, (4, 4, 4)) 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 image_mask = sys.modules["server.jobs.image_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.jobs.image_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): 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_cutout_images_crop_color_and_preserve_source_alpha(self): import json import numpy as np from PIL import Image mesh_job = sys.modules["server.jobs.mesh"] 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, 4, 3] = 64 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 = ( mesh_job._write_cutout_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, 2]), (10, 80, 160, 64)) 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, 64)) def test_mesh_unrefined_cutout_keeps_selection_outside_the_source(self): import json import numpy as np from PIL import Image mesh_job = sys.modules['server.jobs.mesh'] 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 = ( mesh_job._write_unbounded_selection( source_path, polygon, full_selection_path, ) ) color_path, mask_path, analysis_path, bounds_path = ( mesh_job._write_cutout_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()), [-2, 1, 10, 5]) self.assertEqual(color.shape, (4, 12, 4)) self.assertEqual(tuple(color[1, 0]), (0, 0, 0, 0)) self.assertEqual(tuple(color[1, 2]), (10, 80, 160, 255)) self.assertGreater(int(mask[1, 0, 3]), 127) self.assertEqual(tuple(analysis[1, 0]), (0, 0, 0, 0)) def test_mesh_job_keeps_unrefined_outside_selection_with_normal(self): import json 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, 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", "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, { "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()), [-2, 1, 10, 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]), (0, 0, 0, 0)) self.assertEqual(tuple(color_pixels[2, 2, :3]), (20, 40, 60)) self.assertGreater(int(color_pixels[2, 2, 3]), 240) self.assertEqual(tuple(color_pixels[2, 0]), (0, 0, 0, 0)) self.assertEqual(color_pixels.shape[:2], (4, 12)) self.assertEqual(analysis_pixels.shape[:2], (4, 12)) self.assertEqual(normal_pixels.shape[:2], (4, 12)) self.assertEqual(tuple(normal_pixels[0, 0]), (128, 128, 255)) self.assertEqual(tuple(normal_pixels[2, 2]), (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, {"selection": "selection.png", "bounds": "bounds.json"}, ) write_normal.assert_not_called() self.assertFalse((job_directory / "normal.png").exists()) def test_mesh_depth_balloon_uses_cropped_analysis_alpha(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" pixels = np.full((6, 6, 4), (20, 40, 60, 0), dtype=np.uint8) pixels[2:5, 2:5, 3] = 255 Image.fromarray(pixels, mode="RGBA").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, 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]), (128, 128, 255)) 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 balloon = sys.modules["server.jobs.balloon"] source_path = self.storage_root / "source.png" pixels = np.full((6, 6, 4), (20, 40, 60, 255), dtype=np.uint8) Image.fromarray(pixels, mode="RGBA").save(source_path) job_directory = self.storage_root / "job" context = SimpleNamespace( directory=job_directory, progress=lambda *_args: None, check_cancelled=lambda: None, resource=lambda _name: SimpleNamespace( directory=lambda _model: self.storage_root / "model" ), ) received = [] def infer(input_path, *_args, **_kwargs): image = Image.open(input_path).convert("RGBA") received.append(np.asarray(image)) return image, 0.0, 1.0 with patch.object(balloon, "infer_foreground", side_effect=infer): result = balloon.run( context, { "input": str(source_path), "polygon": ((1, 1), (5, 1), (1, 5)), "contour_analysis": True, "depth_analysis": False, "contour_model": "BEN2_BASE", "device": "cpu", }, ) self.assertEqual(result, {"contour": "contour.json"}) self.assertEqual(tuple(received[0][3, 3]), (0, 0, 0, 0)) self.assertTrue((job_directory / "lasso-input.png").is_file()) self.assertTrue((job_directory / "foreground.png").is_file()) self.assertTrue((job_directory / "contour.json").is_file()) def test_balloon_depth_analysis_normalizes_moge_distance(self): import numpy as np balloon = sys.modules["server.jobs.balloon"] y, x = np.indices((7, 9), dtype=np.float32) depth = 2.0 + x * 0.1 + y * 0.2 depth[3, 4] -= 0.5 depth[1, 1] += 3.0 model_mask = np.ones_like(depth, dtype=bool) selection_mask = np.ones_like(depth, dtype=bool) selection_mask[0, 0] = False depth_signal, valid, reference, baseline_vector = balloon.normalized_moge_depth( depth, model_mask, selection_mask, ) normalized_x = x / (depth.shape[1] - 1) * 2.0 - 1.0 normalized_y = y / (depth.shape[0] - 1) * 2.0 - 1.0 fitted_plane = ( baseline_vector[0] * normalized_x + baseline_vector[1] * normalized_y + baseline_vector[2] ) displacement = depth_signal - fitted_plane self.assertFalse(valid[0, 0]) self.assertEqual(depth_signal[0, 0], 0.0) self.assertGreater(float(displacement[3, 4]), 1.0) self.assertLess(float(displacement[3, 4]), 1.01) self.assertLess(displacement[1, 1], -1.0) self.assertGreater(reference, 2.0) self.assertEqual(baseline_vector.shape, (3,)) self.assertLess(float(displacement.min()), -1.0) self.assertLess(float(displacement.max()), 1.01) self.assertLess(float(np.percentile(displacement, 90.0)), 0.01) positive = displacement[valid & (displacement > 0.0)] self.assertAlmostEqual( float(np.sum(np.square(positive)) / np.sum(positive)), 1.0, places=5, ) def test_balloon_relative_depth_ignores_absolute_distance_and_scale(self): import numpy as np balloon = sys.modules["server.jobs.balloon"] y, x = np.indices((8, 10), dtype=np.float32) depth = 4.0 + x * 0.04 + y * 0.03 depth[2:6, 3:7] -= np.asarray( ( (0.1, 0.2, 0.3, 0.1), (0.2, 0.4, 0.8, 0.2), (0.15, 0.3, 0.5, 0.15), (0.05, 0.1, 0.2, 0.05), ), dtype=np.float32, ) mask = np.ones_like(depth, dtype=bool) original = balloon.normalized_moge_depth(depth, mask, mask) farther = balloon.normalized_moge_depth(depth + 100.0, mask, mask) rescaled = balloon.normalized_moge_depth(depth * 7.0, mask, mask) np.testing.assert_allclose(farther[0], original[0], atol=2e-4) 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 balloon = sys.modules["server.jobs.balloon"] source_path = self.storage_root / "depth-source.png" pixels = np.full((6, 6, 4), (80, 100, 120, 255), dtype=np.uint8) Image.fromarray(pixels, mode="RGBA").save(source_path) job_directory = self.storage_root / "depth-job" context = SimpleNamespace( directory=job_directory, progress=lambda *_args: None, check_cancelled=lambda: None, resource=lambda _name: SimpleNamespace(), ) received = [] def infer(_manager, _parameters, images, **_kwargs): image = Image.open(images[0]).convert("RGBA") received.append(np.asarray(image)) height, width = image.size[1], image.size[0] depth = np.linspace( 1.0, 2.0, height * width, dtype=np.float32, ).reshape(height, width) return SimpleNamespace( depth=depth[None], mask=np.ones((1, height, width), dtype=bool), normal=np.broadcast_to( np.asarray((0.0, 0.0, -1.0), dtype=np.float32), (1, height, width, 3), ).copy(), ) with ( patch.object(balloon.moge2, "infer", side_effect=infer), patch.object( balloon, "infer_foreground", side_effect=AssertionError("BEN2 should not run"), ), ): result = balloon.run( context, { "input": str(source_path), "polygon": ((1, 1), (5, 1), (1, 5)), "contour_analysis": False, "depth_analysis": True, "depth_model": "MOGE2_VITS_NORMAL", "resolution_level": 5, "device": "cpu", }, ) self.assertEqual( result, {"depth": "depth.npz", "normal": "normal.png"}, ) self.assertEqual(tuple(received[0][3, 3]), (0, 0, 0, 0)) self.assertFalse((job_directory / "foreground.png").exists()) with np.load(job_directory / "depth.npz", allow_pickle=False) as data: self.assertEqual(data["depth_signal"].shape, (4, 4)) self.assertEqual(data["baseline_vector"].shape, (3,)) self.assertEqual(tuple(data["origin"]), (1, 1)) self.assertEqual(tuple(data["source_size"]), (6, 6)) self.assertTrue(data["valid"].any()) normal = np.asarray(Image.open(job_directory / "normal.png").convert("RGB")) self.assertEqual(normal.shape, (6, 6, 3)) self.assertEqual(tuple(normal[0, 0]), (128, 128, 255)) def test_balloon_contour_uses_alpha_above_half_and_largest_component(self): import numpy as np balloon = sys.modules["server.jobs.balloon"] alpha = np.zeros((6, 7), dtype=np.uint8) alpha[0, 0] = 255 alpha[1:4, 2:5] = 128 alpha[1, 1] = 127 selection_mask = np.ones_like(alpha, dtype=bool) contour = balloon.contour_from_alpha( alpha, selection_mask, origin=(10, 20), ) self.assertEqual( set(map(tuple, contour)), {(12, 21), (15, 21), (15, 24), (12, 24)}, ) if __name__ == "__main__": unittest.main()