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

120
tests/test_addon.py
View File

@@ -1279,7 +1279,7 @@ class AddonTest(unittest.TestCase):
self.assertEqual(paths["status"].parent, paths["result"])
self.assertEqual(paths["status"].stem, paths["result"].name)
def test_new_job_paths_uses_type_time_folder_name(self):
def test_new_job_paths_uses_unique_action_object_names(self):
runtime = importlib.import_module("snapform.runtime")
with TemporaryDirectory() as directory:
jobs_root = Path(directory)
@@ -1300,15 +1300,41 @@ class AddonTest(unittest.TestCase):
self.assertRegex(
depth_paths["result"].name,
r"^depth-plane_\d{8}-\d{6}$",
r"^generate-depth-plane_\d{8}-\d{6}-\d{6}$",
)
self.assertTrue(
background_paths["result"].name.startswith(
"remove-background_"
"remove-image-background_"
)
)
self.assertTrue(
frames_paths["result"].name.startswith("depth-plane_")
frames_paths["result"].name.startswith(
"extract-original-frames_"
)
)
def test_generated_inputs_are_moved_into_job_root(self):
runtime = importlib.import_module("snapform.runtime")
with TemporaryDirectory() as directory:
root = Path(directory)
staging = root / "staging"
staging.mkdir()
(staging / "frame.0001.jpg").write_bytes(b"first")
(staging / "frame.0002.jpg").write_bytes(b"second")
job_directory = root / "job"
job_directory.mkdir()
result = runtime.move_input_into_job(staging, job_directory)
self.assertEqual(result, job_directory)
self.assertFalse(staging.exists())
self.assertEqual(
(job_directory / "frame.0001.jpg").read_bytes(),
b"first",
)
self.assertEqual(
(job_directory / "frame.0002.jpg").read_bytes(),
b"second",
)
def test_empty_image_start_frame_honors_image_user(self):
@@ -1468,9 +1494,83 @@ class AddonTest(unittest.TestCase):
self.assertEqual(
identifiers,
["glb", "gs_ply", "depth_vis", "mini_npz"],
["glb", "gs_ply", "depth_vis", "mini_npz", "mesh"],
)
self.assertEqual(export_format["default"], "glb")
self.assertEqual(export_format["name"], "Output Type")
self.assertEqual(
[item[1] for item in export_format["items"]],
[
"Point Cloud (GLB)",
"Gaussian Splats (PLY)",
"Depth Preview (Images)",
"Depth Data (NPZ)",
"Textured Mesh (GLB)",
],
)
def test_mesh_is_a_da3_output_type(self):
settings = SimpleNamespace(
input_path="",
export_format="mesh",
model="BASE",
gs_model="GIANT_1_1",
process_res=504,
process_res_method="upper_bound_resize",
conf_thresh_percentile=40.0,
num_max_points=1_000_000,
show_cameras=True,
gs_views_interval=1,
video_frames=16,
device="CPU",
use_ray_pose=False,
ref_view_strategy="saddle_balanced",
)
operator = self.addon.RunJob()
operator.job_type = self.addon.operators.RECONSTRUCTION_JOB
with TemporaryDirectory() as directory:
input_path = Path(directory) / "input.png"
input_path.touch()
settings.input_path = str(input_path)
paths = {
"log": Path(directory) / "job.log",
"status": Path(directory) / "job.json",
"result": Path(directory) / "result",
}
with patch.object(
self.addon.runtime,
"new_job_paths",
return_value=paths,
), patch.object(
self.addon.runtime,
"runtime_ready",
return_value=True,
), patch.object(
self.addon.runtime,
"model_ready",
return_value=True,
), patch.object(
self.addon.runtime,
"reconstruct_request",
return_value=self.addon.runtime.HostJobRequest(
{"job_id": "test"}
),
) as reconstruct:
operator._prepare_command(None, settings)
self.assertEqual(reconstruct.call_args.args[1], "mesh")
self.assertEqual(reconstruct.call_args.kwargs["model_key"], "BASE")
def test_confidence_threshold_is_a_bounded_percentage(self):
confidence = self.addon.SnapFormSettings.__annotations__[
"conf_thresh_percentile"
]
self.assertEqual(confidence["name"], "Confidence Threshold")
self.assertEqual(confidence["min"], 0.0)
self.assertEqual(confidence["max"], 100.0)
self.assertEqual(confidence["subtype"], "PERCENTAGE")
def test_gs_model_property_only_includes_gaussian_models(self):
gs_model = self.addon.SnapFormSettings.__annotations__[
@@ -1586,9 +1686,17 @@ class AddonTest(unittest.TestCase):
"DEFAULT_CLOSED",
self.addon.DebugPanel.bl_options,
)
for panel in (self.addon.Da3DebugPanel, self.addon.Ben2DebugPanel):
panels = (
self.addon.Da3OutputDebugPanel,
self.addon.Ben2RemovalDebugPanel,
)
for panel in panels:
self.assertEqual(panel.bl_parent_id, self.addon.DebugPanel.bl_idname)
self.assertIn("DEFAULT_CLOSED", panel.bl_options)
self.assertEqual(
[panel.bl_label for panel in panels],
["DA3 Output Test", "BEN2 Removal Test"],
)
def test_host_panel_exposes_shared_host_state_and_controls(self):
self.assertIn(self.addon.HostPanel, self.addon.CLASSES)

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()