icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Add surface generation and media handling

Browse Source 
- Implemented surface generation functionality in `surface.py`, including geometry estimation and texture writing.
- Introduced media input handling in `input.py`, supporting image and video frame extraction.
- Added model management for geometry adapters in `models/__init__.py`.
- Created BEN2 and DA3 model inference logic in `ben2.py` and `da3.py`, respectively.
- Developed ONNX model inference utilities for BEN2 and DA3 in `onnx_ben2.py` and `onnx_da3.py`.
- Added support for MoGe-2 model inference in `moge2.py` with necessary ONNX utilities.
- Implemented output handling for depth and surface data in `depth.py` and `surface.py`.
- Created debug output utilities for DA3 and MoGe-2 predictions in `debug_da3.py` and `debug_moge2.py`.
- Enhanced ONNX runtime session management in `onnx_runtime.py`.
- Updated `__init__.py` files to include new modules and functions for better organization.
This commit is contained in:
马南
2026-08-18 23:27:10 +08:00
parent 9dfa49602c
commit 835225d264
62 changed files with 2126 additions and 2318 deletions
Download Patch File Download Diff File Expand all files Collapse all files

555
tests/test_worker.py
View File

@@ -74,7 +74,7 @@ class FakeModels:
return self.model_directory
def get_ben2(self, _model_directory, _device):
from server.onnx_ben2 import create_session
from server.models.onnx_ben2 import create_session
return create_session(self.model_directory, "cpu"), 0.0
@@ -108,31 +108,36 @@ class FakeJobContext:
class WorkerTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.media = load_server_module("media")
cls.media = load_server_module("media.input")
cls.geometry = load_server_module("geometry")
cls.background = load_server_module("background")
cls.moge2 = load_server_module("moge2")
cls.da3 = load_server_module("da3")
cls.surface = load_server_module("geometry.surface")
cls.background = load_server_module("jobs.background")
cls.ben2 = load_server_module("models.ben2")
cls.da3_model = load_server_module("models.da3")
cls.moge2_model = load_server_module("models.moge2")
cls.surface_outputs = load_server_module("outputs.surface")
cls.debug_da3_outputs = load_server_module("outputs.debug_da3")
cls.debug_moge2_outputs = load_server_module("outputs.debug_moge2")
cls.depth_job = load_server_module("jobs.depth")
cls.surface_job = load_server_module("jobs.surface")
cls.debug_da3_job = load_server_module("jobs.debug_da3")
cls.debug_moge2_job = load_server_module("jobs.debug_moge2")
cls.debug_ben2_job = load_server_module("jobs.debug_ben2")
cls.frames_job = load_server_module("jobs.frames")
cls.foreground_job = load_server_module("jobs.foreground")
cls.download = load_server_module("model_download")
def test_foreground_preserves_source_rgb_and_ben2_alpha(self):
alpha = SimpleNamespace(size=(2, 1))
source_rgba = SimpleNamespace(size=(2, 1), applied_alpha=None)
source_rgba.putalpha = lambda value: setattr(
source_rgba,
"applied_alpha",
value,
)
source = SimpleNamespace(convert=lambda _mode: source_rgba)
foreground_rgba = SimpleNamespace(
getchannel=lambda _name: alpha,
)
foreground = SimpleNamespace(convert=lambda _mode: foreground_rgba)
def test_surface_has_one_model_independent_builder(self):
self.assertTrue(callable(self.surface.build_surface))
self.assertTrue(callable(self.surface_outputs.write_surface_archive))
self.assertFalse(hasattr(self.surface, "export_da3_surface"))
self.assertFalse(hasattr(self.surface, "export_moge2_surface"))
result = self.background.preserve_source_rgb(source, foreground)
self.assertIs(result, source_rgba)
self.assertIs(result.applied_alpha, alpha)
def test_model_adapters_only_normalize_surface_input(self):
self.assertTrue(callable(self.da3_model.surface_frame))
self.assertTrue(callable(self.moge2_model.surface_frame))
self.assertFalse(hasattr(self.da3_model, "generate_surface"))
self.assertFalse(hasattr(self.moge2_model, "generate_surface"))
def test_collect_images_is_sorted_and_filters_suffixes(self):
with tempfile.TemporaryDirectory() as directory:
@@ -273,70 +278,6 @@ class WorkerTest(unittest.TestCase):
[f"frame.{index:04d}.jpg" for index in range(1, 11)],
)
def test_npz_is_complete_requires_readable_archive(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
valid_path = root / "valid.npz"
np.savez_compressed(
valid_path,
depth=np.zeros((4, 4)),
)
self.assertTrue(self.media.npz_is_complete(valid_path))
empty_path = root / "empty.npz"
empty_path.write_bytes(b"PK\x05\x06" + b"\x00" * 18)
self.assertFalse(self.media.npz_is_complete(empty_path))
def test_wait_for_exported_paths_waits_for_complete_npz(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
npz_path = root / "exports" / "mini_npz" / "results.npz"
npz_path.parent.mkdir(parents=True)
npz_path.touch()
with patch.object(
self.media,
"npz_is_complete",
side_effect=[False, True],
):
result = self.media.wait_for_exported_paths(
root,
"mini_npz",
timeout=2.0,
)
self.assertEqual(result, [npz_path])
def test_move_exported_file_retries_windows_file_lock(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
source = root / "source.npz"
destination = root / "destination.npz"
source.touch()
path_type = type(source)
original_replace = path_type.replace
calls = 0
def locked_once(path, target):
nonlocal calls
calls += 1
if calls == 1:
error = PermissionError("file is in use")
error.winerror = 32
raise error
return original_replace(path, target)
with patch.object(path_type, "replace", new=locked_once), patch.object(
self.media.time,
"sleep",
) as sleep:
result = self.media.move_exported_file(
source,
destination,
)
self.assertEqual(result, destination)
self.assertTrue(destination.is_file())
sleep.assert_called_once_with(0.1)
def test_rejects_unsupported_input_file(self):
with tempfile.TemporaryDirectory() as directory:
input_path = Path(directory) / "input.txt"
@@ -531,7 +472,7 @@ class WorkerTest(unittest.TestCase):
)
fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
fake_torch.inference_mode = nullcontext
fake_ben2 = ModuleType("server.onnx_ben2")
fake_ben2 = ModuleType("server.models.onnx_ben2")
fake_ben2.create_session = lambda _model_dir, _device: object()
fake_ben2.infer = lambda _model, image: (
inference_options.append((image, {})) or FakeImage()
@@ -543,7 +484,7 @@ class WorkerTest(unittest.TestCase):
sys.modules,
{
"torch": fake_torch,
"server.onnx_ben2": fake_ben2,
"server.models.onnx_ben2": fake_ben2,
"PIL": fake_pil,
},
):
@@ -551,7 +492,7 @@ class WorkerTest(unittest.TestCase):
input_path = root / "input.png"
input_path.touch()
context = FakeJobContext(root / "result", "remove-background")
self.background.remove_background(
self.background.run(
context,
{"input": str(input_path), "model": "BEN2", "device": "cpu"},
)
@@ -615,7 +556,7 @@ class WorkerTest(unittest.TestCase):
)
fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
fake_torch.inference_mode = nullcontext
fake_ben2 = ModuleType("server.onnx_ben2")
fake_ben2 = ModuleType("server.models.onnx_ben2")
fake_ben2.create_session = lambda _model_dir, _device: object()
fake_ben2.infer = lambda _model, image: (
inference_options.append((image, {})) or FakeImage()
@@ -627,7 +568,7 @@ class WorkerTest(unittest.TestCase):
sys.modules,
{
"torch": fake_torch,
"server.onnx_ben2": fake_ben2,
"server.models.onnx_ben2": fake_ben2,
"PIL": fake_pil,
},
):
@@ -637,7 +578,7 @@ class WorkerTest(unittest.TestCase):
(input_directory / "frame-0001.jpg").touch()
(input_directory / "frame-0002.jpg").touch()
context = FakeJobContext(root / "result", "remove-background")
self.background.remove_background(
self.background.run(
context,
{
"input": str(input_directory),
@@ -673,7 +614,7 @@ class WorkerTest(unittest.TestCase):
for index in range(1, 6):
(input_directory / f"frame-{index:04d}.jpg").touch()
context = FakeJobContext(root / "result", "extract-frames")
self.media.extract_frames(
self.frames_job.run(
context,
{
"input": str(input_directory),
@@ -694,62 +635,166 @@ class WorkerTest(unittest.TestCase):
self.assertEqual(context.status["progress"], 1.0)
self.assertIn("3 frame", context.status["message"])
def test_uses_da3_export_parameters_for_formats_and_debug_mesh(self):
inference_options = []
mesh_calls = []
def test_da3_debug_saves_only_direct_api_prediction(self):
prediction = SimpleNamespace(
depth=np.full((1, 2, 2), 2.0, dtype=np.float32),
conf=np.ones((1, 2, 2), dtype=np.float32),
extrinsics=np.eye(4, dtype=np.float32)[None, :3],
intrinsics=np.eye(3, dtype=np.float32)[None],
processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8),
is_metric=0,
sky=None,
scale_factor=None,
)
fake_da3 = ModuleType("server.onnx_da3")
fake_da3.create_session = lambda _model_dir, _device: object()
fake_da3 = ModuleType("server.models.onnx_da3")
calls = []
def fake_infer(_session, images, **options):
inference_options.append((images, options))
calls.append((images, options))
return prediction
fake_da3.infer = fake_infer
fake_runtime = ModuleType("server.onnx_runtime")
fake_runtime = ModuleType("server.models.onnx_runtime")
fake_runtime.provider_label = lambda _session: "CPU"
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
input_path = root / "input.png"
input_path.touch()
with patch.dict(
sys.modules,
{
"server.models.onnx_da3": fake_da3,
"server.models.onnx_runtime": fake_runtime,
},
):
context = FakeJobContext(root / "result", "debug-da3")
self.debug_da3_job.run(
context,
{
"input": str(input_path),
"model": "DA3",
"video_frames": 1,
"device": "cpu",
"process_res": 518,
"process_res_method": "lower_bound_resize",
"ref_view_strategy": "middle",
},
)
with np.load(context.directory / "prediction.npz") as result:
self.assertEqual(
set(result.files),
{
"depth",
"conf",
"extrinsics",
"intrinsics",
"processed_images",
"is_metric",
},
)
np.testing.assert_array_equal(result["conf"], prediction.conf)
self.assertFalse(any(context.directory.glob("*.glb")))
self.assertFalse(any(context.directory.glob("*.png")))
self.assertEqual(calls[0][0], [str(input_path)])
self.assertEqual(calls[0][1]["process_res"], 518)
def test_moge2_debug_saves_each_direct_api_dictionary(self):
prediction = {
"points": np.ones((2, 2, 3), dtype=np.float32),
"depth": np.ones((2, 2), dtype=np.float32),
"normal": np.ones((2, 2, 3), dtype=np.float32),
"mask": np.ones((2, 2), dtype=bool),
"intrinsics": np.eye(3, dtype=np.float32),
}
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
frames = root / "frames"
frames.mkdir()
(frames / "a.png").touch()
(frames / "b.png").touch()
context = FakeJobContext(root / "result", "debug-moge2")
with patch.object(
self.debug_moge2_job,
"infer_one",
return_value=prediction,
) as infer_one:
self.debug_moge2_job.run(
context,
{
"input": str(frames),
"model": "MOGE2_VITS_NORMAL",
"video_frames": 2,
"device": "cpu",
"resolution_level": 5,
},
)
self.assertEqual(infer_one.call_count, 2)
archives = sorted(context.directory.glob("prediction.*.npz"))
self.assertEqual(len(archives), 2)
with np.load(archives[0]) as result:
self.assertEqual(set(result.files), set(prediction))
self.assertFalse(any(context.directory.glob("*.glb")))
def test_ben2_debug_saves_direct_api_foreground(self):
class Foreground:
def save(self, path):
Path(path).write_bytes(b"foreground")
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
input_path = root / "input.png"
input_path.touch()
context = FakeJobContext(root / "result", "debug-ben2")
with patch.object(
self.debug_ben2_job,
"infer_foreground",
return_value=(Foreground(), 1.0, 2.0),
) as infer_foreground:
self.debug_ben2_job.run(
context,
{
"input": str(input_path),
"model": "BEN2_BASE",
"video_frames": 1,
"device": "cpu",
},
)
infer_foreground.assert_called_once()
self.assertEqual(
(context.directory / "foreground.png").read_bytes(),
b"foreground",
)
def test_da3_removes_background_after_depth_using_original_input(self):
order = []
inference_inputs = []
fake_da3 = ModuleType("server.models.onnx_da3")
def fake_infer(_session, images, **_options):
order.append("depth")
inference_inputs.append(images)
return SimpleNamespace(
depth=np.full((1, 2, 2), 2.0, dtype=np.float32),
conf=np.ones((1, 2, 2), dtype=np.float32),
extrinsics=np.eye(4, dtype=np.float32)[None, :3],
intrinsics=np.eye(3, dtype=np.float32)[None],
processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8),
)
fake_da3.infer = fake_infer
fake_runtime = ModuleType("server.models.onnx_runtime")
fake_runtime.provider_label = lambda _session: "CPU"
def fake_mesh_export(
prediction,
images,
output_directory,
confidence_percentile,
filter_edges,
max_edge_length,
):
mesh_calls.append(
(
prediction,
images,
confidence_percentile,
filter_edges,
max_edge_length,
)
)
output_path = Path(output_directory) / "mesh.glb"
output_path.touch()
return output_path
class FakeForeground:
def convert(self, _mode):
return self
def fake_point_cloud(_prediction, output_directory, *_args):
(Path(output_directory) / "scene.glb").touch()
def save(self, path, **_options):
Path(path).touch()
def fake_depth_preview(_prediction, output_directory):
target = Path(output_directory) / "depth_vis"
target.mkdir()
(target / "0000.jpg").touch()
def fake_mini_npz(_prediction, output_directory):
np.savez_compressed(
Path(output_directory) / "results.npz",
depth=np.zeros((1, 2, 2)),
)
def fake_foreground(image, *_args, **_kwargs):
order.append("background")
self.assertEqual(image, input_path)
return FakeForeground(), 1.0, 2.0
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
@@ -758,77 +803,40 @@ class WorkerTest(unittest.TestCase):
with patch.dict(
sys.modules,
{
"server.onnx_da3": fake_da3,
"server.onnx_runtime": fake_runtime,
"server.models.onnx_da3": fake_da3,
"server.models.onnx_runtime": fake_runtime,
},
), patch.object(
self.da3,
"export_textured_mesh_glb",
side_effect=fake_mesh_export,
), patch.object(
self.da3,
"export_da3_point_cloud",
side_effect=fake_point_cloud,
), patch.object(
self.da3,
"export_da3_depth_preview",
side_effect=fake_depth_preview,
), patch.object(
self.da3,
"export_da3_mini_npz",
side_effect=fake_mini_npz,
self.foreground_job,
"infer_foreground",
side_effect=fake_foreground,
):
debug_formats = (
"glb",
"depth_vis",
"mini_npz",
"mesh",
context = FakeJobContext(
root / "result",
"generate-depth",
)
self.depth_job.run(
context,
{
"input": str(input_path),
"model": "DA3",
"remove_background_model": "BEN2",
"video_frames": 1,
"export_format": "mini_npz",
"device": "cpu",
"process_res": 518,
"process_res_method": "upper_bound_resize",
"ref_view_strategy": "first",
"inference_mode": "per_frame",
},
)
for export_format in debug_formats:
context = FakeJobContext(
root / export_format,
"debug-da3",
)
self.da3.generate_da3(
context,
{
"input": str(input_path),
"model": "DA3",
"video_frames": 16,
"export_format": export_format,
"device": "cpu",
"process_res": 518,
"process_res_method": "lower_bound_resize",
"ref_view_strategy": "middle",
"conf_thresh_percentile": 35.0,
"num_max_points": 12345,
"mesh_filter_edges": True,
"mesh_max_edge_length": 0.2,
},
)
self.assertEqual(context.status["progress"], 1.0)
self.assertIn(export_format, context.status["message"])
result_directory = context.directory
if export_format == "depth_vis":
self.assertTrue((result_directory / "0000.jpg").is_file())
self.assertFalse((result_directory / "depth_vis").exists())
self.assertEqual(len(inference_options), 4)
self.assertEqual(len(mesh_calls), 1)
for images, options in inference_options:
self.assertEqual(images, [str(input_path)])
self.assertEqual(options["process_res"], 518)
self.assertEqual(
options["process_res_method"],
"lower_bound_resize",
)
self.assertEqual(options["ref_view_strategy"], "middle")
self.assertEqual(mesh_calls[0][1], [input_path])
self.assertEqual(mesh_calls[0][2:], (35.0, True, 0.2))
self.assertEqual(order, ["depth", "background"])
self.assertEqual(inference_inputs, [[str(input_path)]])
def test_reconstruct_copies_frames_for_multi_frame_input(self):
inference_inputs = []
fake_da3 = ModuleType("server.onnx_da3")
fake_da3 = ModuleType("server.models.onnx_da3")
fake_da3.create_session = lambda _model_dir, _device: object()
def fake_infer(_session, images, **_options):
@@ -842,7 +850,7 @@ class WorkerTest(unittest.TestCase):
)
fake_da3.infer = fake_infer
fake_runtime = ModuleType("server.onnx_runtime")
fake_runtime = ModuleType("server.models.onnx_runtime")
fake_runtime.provider_label = lambda _session: "CPU"
with tempfile.TemporaryDirectory() as directory:
@@ -854,15 +862,15 @@ class WorkerTest(unittest.TestCase):
with patch.dict(
sys.modules,
{
"server.onnx_da3": fake_da3,
"server.onnx_runtime": fake_runtime,
"server.models.onnx_da3": fake_da3,
"server.models.onnx_runtime": fake_runtime,
},
):
context = FakeJobContext(
root / "result",
"generate-da3-depth",
"generate-depth",
)
self.da3.generate_da3(
self.depth_job.run(
context,
{
"input": str(input_directory),
@@ -1099,68 +1107,51 @@ class WorkerTest(unittest.TestCase):
)
np.testing.assert_array_equal(compact_faces, [[0, 1, 2]])
def test_surface_background_faces_are_removed_before_export(self):
def test_surface_background_keeps_faces_with_any_foreground_vertex(self):
faces = np.asarray(
((0, 1, 4), (0, 4, 3), (1, 2, 5), (1, 5, 4)),
dtype=np.int32,
)
foreground = np.asarray(
((False, True, True), (False, True, True)),
((False, False, False), (False, True, False)),
dtype=bool,
)
filtered = self.geometry.remove_background_faces(faces, foreground)
np.testing.assert_array_equal(filtered, faces[2:])
np.testing.assert_array_equal(filtered, faces[[0, 1, 3]])
def test_surface_background_mask_keeps_any_nonzero_alpha(self):
from PIL import Image
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
for alpha, expected in ((0, False), (1, True), (127, True)):
path = root / f"alpha-{alpha}.png"
Image.new("RGBA", (2, 2), (10, 20, 30, alpha)).save(path)
mask = self.geometry.foreground_mesh_mask(path, (2, 2))
self.assertEqual(mask.dtype, np.bool_)
self.assertTrue(np.all(mask == expected))
def test_image_mesh_data_keeps_camera_geometry_and_uv(self):
class FakeImage:
def __enter__(self):
return self
def __exit__(self, *_args):
return None
def convert(self, mode):
self.mode = mode
return self
def save(self, path, **_options):
Path(path).write_bytes(b"texture")
fake_pil = ModuleType("PIL")
fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage())
depth = np.full((1, 3, 3), 2.0, dtype=np.float32)
confidence = np.ones((1, 3, 3), dtype=np.float32)
depth = np.full((3, 3), 2.0, dtype=np.float32)
confidence = np.ones((3, 3), dtype=np.float32)
intrinsics = np.asarray(
[[[2.0, 0.0, 1.0], [0.0, 2.0, 1.0], [0.0, 0.0, 1.0]]],
[[2.0, 0.0, 1.0], [0.0, 2.0, 1.0], [0.0, 0.0, 1.0]],
dtype=np.float32,
)
prediction = SimpleNamespace(
frame = self.surface.GeometryFrame(
depth=depth,
conf=confidence,
points=self.geometry.unproject_camera_depth(depth, intrinsics),
intrinsics=intrinsics,
valid_mask=np.ones(depth.shape, dtype=bool),
confidence=confidence,
)
progress = []
with tempfile.TemporaryDirectory() as directory, patch.dict(
sys.modules,
{"PIL": fake_pil},
), patch.object(
self.geometry,
"compact_mesh",
side_effect=lambda vertices, faces: (
vertices,
faces,
),
):
result_path, texture_path = self.geometry.export_image_mesh_data(
prediction,
Path(directory) / "input.png",
Path(directory) / "result",
progress_callback=lambda step, total, label: progress.append(
(step, total, label)
),
mesh = self.surface.build_surface(frame)
with tempfile.TemporaryDirectory() as directory:
result_path = self.surface_outputs.write_surface_archive(
mesh,
Path(directory) / "mesh.npz",
)
with np.load(result_path) as result:
self.assertEqual(result["vertices"].shape, (9, 3))
@@ -1170,7 +1161,7 @@ class WorkerTest(unittest.TestCase):
self.assertNotIn("postprocess_method", result.files)
self.assertEqual(result["point_depth"].shape, (9,))
self.assertEqual(
result["face_confidence_percentile"].shape,
result["face_confidence"].shape,
(8,),
)
self.assertEqual(
@@ -1183,17 +1174,67 @@ class WorkerTest(unittest.TestCase):
"image_size",
"reference_depth",
"point_depth",
"face_confidence_percentile",
"face_confidence",
"face_depth_discontinuity",
},
)
np.testing.assert_allclose(result["uv"][0], [0.0, 1.0])
np.testing.assert_allclose(result["uv"][-1], [1.0, 0.0])
self.assertTrue(texture_path.is_file())
self.assertEqual([item[0] for item in progress], list(range(1, 7)))
self.assertTrue(all(item[1] == 6 for item in progress))
self.assertNotIn("normal_map", result.files)
def test_prediction_attributes_keep_debug_structure(self):
def test_moge_surface_uses_the_same_archive_without_normal_data(self):
depth = np.full((3, 3), 2.0, dtype=np.float32)
intrinsics = np.asarray(
[[2.0, 0.0, 1.0], [0.0, 2.0, 1.0], [0.0, 0.0, 1.0]],
dtype=np.float32,
)
frame = self.surface.GeometryFrame(
depth=depth,
points=self.geometry.unproject_camera_depth(depth, intrinsics),
intrinsics=intrinsics,
valid_mask=np.ones(depth.shape, dtype=bool),
model_mask=np.asarray(
(
(False, False, False),
(True, True, True),
(True, True, True),
),
dtype=bool,
),
)
mesh = self.surface.build_surface(frame)
with tempfile.TemporaryDirectory() as directory:
path = self.surface_outputs.write_surface_archive(
mesh,
Path(directory) / "mesh.npz",
)
with np.load(path) as result:
self.assertIn("face_mask", result.files)
self.assertNotIn("face_confidence", result.files)
self.assertNotIn("normal", result.files)
self.assertNotIn("normal_map", result.files)
def test_moge_surface_normal_is_exported_as_a_texture(self):
from PIL import Image
normal = np.asarray(
[[[0.0, 0.0, -1.0]]],
dtype=np.float32,
)
with tempfile.TemporaryDirectory() as directory:
path = self.surface_outputs.write_normal_texture(
normal,
Path(directory) / "normal.png",
)
with Image.open(path) as image:
self.assertEqual(image.mode, "RGB")
self.assertEqual(image.size, (1, 1))
red, green, blue = image.getpixel((0, 0))
self.assertIn(red, {127, 128})
self.assertIn(green, {127, 128})
self.assertEqual(blue, 255)
def test_da3_debug_writer_keeps_api_values_without_derived_attributes(self):
prediction = SimpleNamespace(
depth=np.ones((1, 2, 3), dtype=np.float32),
conf=np.arange(6, dtype=np.float32).reshape(1, 2, 3),
@@ -1202,18 +1243,18 @@ class WorkerTest(unittest.TestCase):
)
with tempfile.TemporaryDirectory() as directory:
output_path = self.geometry.export_prediction_attributes(
output_path = self.debug_da3_outputs.write_prediction(
prediction,
directory,
Path(directory) / "prediction.npz",
)
with np.load(output_path) as attributes:
self.assertEqual(attributes["depth"].shape, (1, 2, 3))
self.assertEqual(attributes["conf"].shape, (1, 2, 3))
self.assertEqual(
attributes["confidence_percentile"].shape,
(1, 2, 3),
)
self.assertNotIn("confidence_percentile", attributes.files)
if __name__ == "__main__":
unittest.main()