Implement Balloon Depth Sampling and Validation Tools

- Added depth.py for MoGe-2 depth sampling and balloon deformation functionality.
- Introduced validate_balloon_asset.py to validate balloon shape modifiers and material properties in Blender.
- Enhanced test_server.py with new unit tests for balloon depth analysis and job progress tracking.
- Updated existing tests to improve code clarity and maintainability.

tests/test_server.py

@@ -15,7 +15,8 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1]
 def load_server_module():
     server_directory = PROJECT_ROOT / "src" / "snapform" / "server"
     specification = importlib.util.spec_from_file_location(
-        "server", server_directory / "__init__.py",
+        "server",
+        server_directory / "__init__.py",
         submodule_search_locations=[str(server_directory)],
     )
     module = importlib.util.module_from_spec(specification)
@@ -51,14 +52,17 @@ class ServerTest(unittest.TestCase):
         manager = self.server_module.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),
+        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,
@@ -215,7 +219,7 @@ class ServerTest(unittest.TestCase):
             {
                 "download-model",
                 "download-required-models",
-                "refine-balloon",
+                "analyze-balloon",
                 "remove-background",
                 "extract-frames",
                 "generate-depth",
@@ -254,15 +258,11 @@ class ServerTest(unittest.TestCase):
 
     def test_cancelling_job_keeps_server_and_model_cache_alive(self):
         fake_models = SimpleNamespace(
-            snapshot=lambda: {
-                "loaded": {"ben2": "BEN2_BASE", "da3": ""}
-            },
+            snapshot=lambda: {"loaded": {"ben2": "BEN2_BASE", "da3": ""}},
             clear=lambda: None,
             close=lambda: None,
         )
-        with patch.object(
-            self.server_module, "ModelManager", return_value=fake_models
-        ):
+        with patch.object(self.server_module, "ModelManager", return_value=fake_models):
             server = self.create_server()
         started = threading.Event()
 
@@ -287,9 +287,7 @@ class ServerTest(unittest.TestCase):
         )
 
     def test_job_progress_is_kept_in_server_memory(self):
-        server = self.server_module.JobServer(
-            "Test", storage_root=self.storage_root
-        )
+        server = self.server_module.JobServer("Test", storage_root=self.storage_root)
 
         @server.job("example")
         def example(context, _parameters):
@@ -326,9 +324,7 @@ class ServerTest(unittest.TestCase):
             clear=lambda: events.append("clear"),
             close=lambda: events.append("close"),
         )
-        server = self.server_module.JobServer(
-            "Test", storage_root=self.storage_root
-        )
+        server = self.server_module.JobServer("Test", storage_root=self.storage_root)
         server.add_resource("model_manager", resource)
 
         @server.job("example")
@@ -356,9 +352,7 @@ class ServerTest(unittest.TestCase):
             clear=lambda: None,
             close=lambda: None,
         )
-        with patch.object(
-            self.server_module, "ModelManager", return_value=fake_models
-        ):
+        with patch.object(self.server_module, "ModelManager", return_value=fake_models):
             server = self.create_server()
         with tempfile.TemporaryDirectory() as directory:
             root = Path(directory)
@@ -376,9 +370,7 @@ class ServerTest(unittest.TestCase):
                 job_id="job",
             )
             status = self.wait_for_job(context)
-            self.assertTrue(
-                (context.directory / "color.0001.jpg").is_file()
-            )
+            self.assertTrue((context.directory / "color.0001.jpg").is_file())
 
         self.assertEqual(status["state"], "succeeded")
 
@@ -395,7 +387,7 @@ class ServerTest(unittest.TestCase):
             pixels[1, 1, 3] = 64
             Image.fromarray(pixels, mode="RGBA").save(source_path)
 
-            _path, _selection, origin = balloon.create_lasso_input(
+            _path, _selection, origin, source_size = balloon.create_lasso_input(
                 source_path,
                 ((1, 1), (5, 1), (1, 5)),
                 output_path,
@@ -403,6 +395,7 @@ class ServerTest(unittest.TestCase):
             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))
@@ -437,7 +430,9 @@ class ServerTest(unittest.TestCase):
                 {
                     "input": str(source_path),
                     "polygon": ((1, 1), (5, 1), (1, 5)),
-                    "model": "BEN2_BASE",
+                    "contour_analysis": True,
+                    "depth_analysis": False,
+                    "contour_model": "BEN2_BASE",
                     "device": "cpu",
                 },
             )
@@ -448,6 +443,148 @@ class ServerTest(unittest.TestCase):
         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
 
@@ -469,5 +606,6 @@ class ServerTest(unittest.TestCase):
             {(12, 21), (15, 21), (15, 24), (12, 24)},
         )
 
+
 if __name__ == "__main__":
     unittest.main()