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

Refactor and enhance DA3 and BEN2 debug panels, introduce mesh export functionality

Browse Source 
- Renamed DA3 and BEN2 debug panels for clarity and consistency.
- Added a new function to draw DA3 input settings to reduce code duplication.
- Implemented mesh export functionality, including confidence thresholding and edge filtering.
- Updated export format options to include new mesh type and adjusted related properties.
- Enhanced job handling to move generated inputs into job root and manage temporary directories.
- Improved error handling and file management during export processes.
- Added unit tests for new mesh functionality and ensured existing tests cover changes.
This commit is contained in:
马南
2026-08-09 21:31:08 +08:00
parent 8aac4822c6
commit 6d4d1c4ec2
10 changed files with 778 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files

158
tests/test_worker.py
View File

@@ -474,6 +474,38 @@ class WorkerTest(unittest.TestCase):
)
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.worker.time,
"sleep",
) as sleep:
result = self.worker.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"
@@ -850,9 +882,10 @@ class WorkerTest(unittest.TestCase):
self.assertEqual(status["progress"], 1.0)
self.assertIn("3 frame", status["message"])
def test_uses_official_da3_export_parameters_for_all_four_formats(self):
def test_uses_da3_export_parameters_for_formats_and_debug_mesh(self):
inference_options = []
export_calls = []
mesh_calls = []
class FakeModel:
@classmethod
@@ -910,6 +943,27 @@ class WorkerTest(unittest.TestCase):
fake_package = ModuleType("depth_anything_3")
fake_package.__path__ = []
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
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
input_path = root / "input.png"
@@ -923,6 +977,10 @@ class WorkerTest(unittest.TestCase):
"depth_anything_3.utils": fake_utils,
"depth_anything_3.utils.export": fake_export_module,
},
), patch.object(
self.worker,
"export_textured_mesh_glb",
side_effect=fake_mesh_export,
):
for export_format in self.worker.EXPORT_FORMATS:
args = SimpleNamespace(
@@ -941,6 +999,8 @@ class WorkerTest(unittest.TestCase):
num_max_points=12345,
show_cameras=0,
gs_views_interval=2,
mesh_filter_edges=True,
mesh_max_edge_length=0.2,
)
self.worker.reconstruct(args)
status = json.loads(
@@ -948,9 +1008,17 @@ class WorkerTest(unittest.TestCase):
)
self.assertEqual(status["progress"], 1.0)
self.assertIn(export_format, status["message"])
result_directory = Path(args.output_dir)
if export_format == "gs_ply":
self.assertTrue((result_directory / "0000.ply").is_file())
self.assertFalse((result_directory / "gs_ply").exists())
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(inference_options), 5)
self.assertEqual(len(export_calls), 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)
@@ -960,8 +1028,15 @@ class WorkerTest(unittest.TestCase):
)
self.assertTrue(options["use_ray_pose"])
self.assertEqual(options["ref_view_strategy"], "middle")
for _prediction, export_format, _export_dir, options in export_calls:
expected_depth = -1.0 if export_format == "mini_npz" else 2.0
for (
_prediction,
export_format,
_export_dir,
options,
) in export_calls:
expected_depth = (
-1.0 if export_format == "mini_npz" else 2.0
)
np.testing.assert_array_equal(
_prediction.depth,
np.full((1, 2, 2), expected_depth, dtype=np.float32),
@@ -985,9 +1060,14 @@ class WorkerTest(unittest.TestCase):
)
else:
self.assertEqual(options, {})
self.assertEqual(mesh_calls[0][1], [input_path])
self.assertEqual(mesh_calls[0][2:], (35.0, True, 0.2))
infer_gs_values = [options["infer_gs"] for _, options in inference_options]
self.assertEqual(infer_gs_values, [False, True, False, False])
self.assertEqual(
infer_gs_values,
[False, True, False, False, False],
)
def test_reconstruct_copies_frames_for_multi_frame_input(self):
class FakeModel:
@@ -1080,6 +1160,74 @@ class WorkerTest(unittest.TestCase):
status = json.loads(Path(args.status).read_text(encoding="utf-8"))
self.assertEqual(status["progress"], 1.0)
def test_mesh_confidence_uses_percentile(self):
depth = np.ones((1, 2, 2), dtype=np.float32)
confidence = np.array([[[1.0, 2.0], [3.0, 4.0]]])
threshold = self.worker.confidence_percentile_threshold(
depth,
confidence,
50.0,
)
self.assertEqual(threshold, 2.5)
def test_mesh_faces_apply_confidence_threshold_and_face_direction(self):
points = np.array(
[
[[0.0, 0.0, 1.0], [1.0, 0.0, 1.0], [2.0, 0.0, 1.0]],
[[0.0, 1.0, 1.0], [1.0, 1.0, 1.0], [2.0, 1.0, 1.0]],
],
dtype=np.float32,
)
depth = np.ones((2, 3), dtype=np.float32)
confidence = np.ones((2, 3), dtype=np.float32)
confidence[:, 0] = 0.0
faces = self.worker.mesh_triangle_faces(
points,
depth,
confidence,
confidence_threshold=0.5,
filter_edges=False,
max_edge_length=0.0,
)
np.testing.assert_array_equal(faces, [[1, 4, 5], [1, 5, 2]])
vertices = points.reshape(-1, 3)[faces[0]]
normal = np.cross(
vertices[1] - vertices[0],
vertices[2] - vertices[0],
)
self.assertLess(normal[2], 0.0)
def test_mesh_percentile_reduces_face_count(self):
depth = np.ones((3, 3), dtype=np.float32)
confidence = np.arange(1, 10, dtype=np.float32).reshape(3, 3)
x, y = np.meshgrid(
np.arange(3, dtype=np.float32),
np.arange(3, dtype=np.float32),
)
points = np.stack((x, y, depth), axis=-1)
face_counts = []
for percentile in (0.0, 50.0, 100.0):
threshold = self.worker.confidence_percentile_threshold(
depth,
confidence,
percentile,
)
faces = self.worker.mesh_triangle_faces(
points,
depth,
confidence,
confidence_threshold=threshold,
filter_edges=False,
max_edge_length=0.0,
)
face_counts.append(len(faces))
self.assertEqual(face_counts, [8, 2, 0])
if __name__ == "__main__":
unittest.main()