import importlib.util import json import sys from types import SimpleNamespace from unittest.mock import patch from server_runtime_test_support import ServerTestCase class GeometryArtifactsTest(ServerTestCase): 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 / "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_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( [[ (127, 127, 255), (255, 127, 127), (127, 0, 127), (128, 128, 255), ]], 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_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_depth_texture( point_frame, self.storage_root / "nan-points.exr", ) depth_frame = GeometryFrame( **base, points=np.asarray([[(0, 0, 1), (0, 0, np.inf)]], dtype=np.float32), ) with self.assertRaisesRegex(ValueError, "Point field"): depth.write_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, generate_depth=True, normal_mode="OBJECT", ) infer.assert_called_once() self.assertEqual( set(result), {"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_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, ) depth_texture = context.directory / "depth.exr" metadata = context.directory / "depth.json" normal = context.directory / "object-normal.png" with patch.object( job, "generate_moge2_artifacts", return_value={ "depth": depth_texture, "depth_metadata": metadata, "object_normal": normal, }, ) as generate: result = job.run( context, { "input": str(input_path), "generate_depth": True, "normal_mode": "OBJECT", }, ) self.assertEqual( result, { "depth": "depth.exr", "depth_metadata": "depth.json", "object_normal": "object-normal.png", }, ) generate.assert_called_once_with( context, { "input": str(input_path), "generate_depth": True, "normal_mode": "OBJECT", }, input_path, generate_depth=True, normal_mode="OBJECT", correct_edge_depth=True, alpha_threshold=0.9, )