AnyImage/tests/test_server.py

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.cutout_depth")
        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_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_depth_balloon(
                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_depth_analysis_normalizes_moge_distance(self):
        import numpy as np

        balloon = sys.modules["server.jobs.cutout_depth"]
        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_depth_plane_is_stable_for_a_slender_diagonal_mask(self):
        import numpy as np

        balloon = sys.modules["server.jobs.cutout_depth"]
        y, x = np.indices((101, 101), dtype=np.float64)
        normalized_x = x / 50.0 - 1.0
        normalized_y = y / 50.0 - 1.0
        distance = y - x
        selection = np.abs(distance) <= 1.0
        along = (normalized_x + normalized_y) * 0.5
        depth = 4.0 + 0.4 * along
        depth += np.where(distance > 0.0, 0.05, 0.0)
        depth -= np.where(distance < 0.0, 0.05, 0.0)

        _signal, valid, _reference, baseline = balloon.normalized_moge_depth(
            depth,
            np.ones_like(selection),
            selection,
        )
        fitted = (
            baseline[0] * normalized_x
            + baseline[1] * normalized_y
            + baseline[2]
        )

        self.assertLessEqual(
            float(np.ptp(fitted)),
            float(np.ptp(fitted[valid])) * 1.25,
        )
        self.assertGreater(
            abs(float(np.corrcoef(fitted[valid], along[valid])[0, 1])),
            0.99,
        )

    def test_balloon_relative_depth_ignores_absolute_distance_and_scale(self):
        import numpy as np

        balloon = sys.modules["server.jobs.cutout_depth"]
        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)

if __name__ == "__main__":
    unittest.main()