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

feat: Implement background depth distance feature for depth processing

Browse Source 
- Added `background_depth_distance` property to `ActiveJob` and `RunJob` classes to manage background depth settings.
- Updated job execution to handle background depth distance, allowing for fixed negative displacement in depth calculations.
- Introduced new functions in the worker module to apply background depth distance to depth frames and manage original background depth.
- Enhanced tests to cover new functionality, ensuring correct behavior of background depth processing and reference depth extraction.
- Refactored job submission to run in a separate thread, improving responsiveness during execution.
This commit is contained in:
马南
2026-08-09 16:34:41 +08:00
parent 2e8af26cdf
commit 8aac4822c6
9 changed files with 725 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files

128
tests/test_addon.py
View File

@@ -1,5 +1,7 @@
import importlib
import sys
import threading
import time
import unittest
from pathlib import Path
from tempfile import TemporaryDirectory
@@ -321,7 +323,7 @@ class AddonTest(unittest.TestCase):
def test_original_depth_background_starts_independent_job(self):
operator = self.addon.ConvertImageToDepthPlane()
operator.snapform_depth_source = "ORIGINAL"
operator.snapform_remove_background = True
operator.snapform_background_depth_distance = 2.0
operator.snapform_displacement_via = "MODIFIER"
operator.snapform_unshaded = False
operator.snapform_subdivision_level = 5
@@ -351,12 +353,12 @@ class AddonTest(unittest.TestCase):
self.assertEqual(result, {"FINISHED"})
self.assertEqual(calls[0]["job_type"], "ORIGINAL_BACKGROUND")
self.assertFalse(calls[0]["snapform_delete_source"])
self.assertEqual(calls[0]["snapform_background_depth_distance"], 2.0)
def test_original_depth_animation_extracts_frames_without_background(self):
operator = self.addon.ConvertImageToDepthPlane()
operator.snapform_depth_source = "ORIGINAL"
operator.snapform_remove_background = False
operator.snapform_background_depth_distance = 0.0
operator.snapform_displacement_via = "MODIFIER"
operator.snapform_unshaded = False
operator.snapform_subdivision_level = 5
@@ -417,9 +419,11 @@ class AddonTest(unittest.TestCase):
self.assertFalse(
operator.__annotations__["snapform_unshaded"]["default"]
)
self.assertFalse(
operator.__annotations__["snapform_remove_background"]["default"]
)
background_offset = operator.__annotations__[
"snapform_background_depth_distance"
]
self.assertEqual(background_offset["default"], 0.0)
self.assertEqual(background_offset["min"], 0.0)
self.assertEqual(
operator.__annotations__["snapform_subdivision_level"]["default"],
5,
@@ -430,7 +434,7 @@ class AddonTest(unittest.TestCase):
"snapform_depth_source",
"snapform_displacement_via",
"snapform_unshaded",
"snapform_remove_background",
"snapform_background_depth_distance",
"snapform_model",
"snapform_use_ray_pose",
"snapform_process_res_factor",
@@ -537,7 +541,7 @@ class AddonTest(unittest.TestCase):
self.assertEqual(result, {"FINISHED"})
self.assertEqual(scene.snapform_settings.video_frames, 24)
def test_conversion_execute_uses_remove_background_frame_range(self):
def test_conversion_execute_uses_background_offset_frame_range(self):
source_object = SimpleNamespace(
name="Reference",
data=SimpleNamespace(
@@ -559,7 +563,7 @@ class AddonTest(unittest.TestCase):
operator.snapform_process_res_factor = 1.0
operator.snapform_displacement_via = "MODIFIER"
operator.snapform_unshaded = False
operator.snapform_remove_background = False
operator.snapform_background_depth_distance = 1.0
operator.snapform_subdivision_level = 5
operator.snapform_delete_original = False
calls = []
@@ -583,6 +587,7 @@ class AddonTest(unittest.TestCase):
self.assertEqual(result, {"FINISHED"})
self.assertEqual(scene.snapform_settings.video_frames, 24)
self.assertEqual(calls[0]["job_type"], "IMAGE_DEPTH")
self.assertEqual(calls[0]["snapform_background_depth_distance"], 1.0)
def test_conversion_dialog_only_lists_installed_models(self):
items_callback = self.addon.ConvertImageToDepthPlane.__annotations__[
@@ -687,7 +692,7 @@ class AddonTest(unittest.TestCase):
operator.snapform_process_res_factor = 0.5
operator.snapform_displacement_via = "MATERIAL"
operator.snapform_unshaded = True
operator.snapform_remove_background = True
operator.snapform_background_depth_distance = 2.5
operator.snapform_subdivision_level = 5
operator.snapform_delete_original = True
image = SimpleNamespace(
@@ -732,7 +737,7 @@ class AddonTest(unittest.TestCase):
"MATERIAL",
)
self.assertTrue(calls[0]["snapform_unshaded"])
self.assertTrue(calls[0]["snapform_remove_background"])
self.assertEqual(calls[0]["snapform_background_depth_distance"], 2.5)
self.assertEqual(calls[0]["snapform_depth_data_name"], "img.depth")
self.assertTrue(calls[0]["snapform_delete_input"])
self.assertEqual(
@@ -944,7 +949,7 @@ class AddonTest(unittest.TestCase):
operator.snapform_model_key = "BASE"
operator.snapform_use_ray_pose = True
operator.snapform_process_res = 1024
operator.snapform_remove_background = True
operator.snapform_background_depth_distance = 1.0
with TemporaryDirectory() as directory:
input_path = Path(directory) / "input.png"
@@ -1002,12 +1007,105 @@ class AddonTest(unittest.TestCase):
Path("status.json"),
Path("result"),
remove_background=True,
background_depth_distance=2.5,
)
self.assertEqual(
request.payload["parameters"]["remove_background_model_dir"],
str(Path("ben2_base")),
)
self.assertEqual(
request.payload["parameters"]["background_depth_distance"],
2.5,
)
def test_host_busy_snapshot_is_released_for_completed_local_job(self):
runtime = self.addon.runtime
previous_snapshot = runtime._host_snapshot
runtime._host_snapshot = {
"state": "BUSY",
"busy": True,
"active_job": {"job_id": "depth-job"},
}
try:
runtime.mark_host_job_complete("depth-job")
self.assertEqual(runtime.host_status()["state"], "READY")
self.assertFalse(runtime.host_busy())
finally:
runtime._host_snapshot = previous_snapshot
def test_job_submission_starts_without_blocking_execute(self):
operator = self.addon.RunJob()
operator.job_type = self.addon.operators.IMAGE_DEPTH_JOB
operator.snapform_source_object_name = "Reference"
operator.snapform_depth_data_name = "img.depth"
operator.snapform_process_res = 1024
operator.snapform_subdivision_level = 5
operator.snapform_delete_source = False
operator.snapform_delete_input = False
operator.snapform_displacement_via = "MODIFIER"
operator.snapform_unshaded = False
operator.snapform_background_depth_distance = 0.0
started = threading.Event()
release = threading.Event()
def submit(_request):
started.set()
release.wait(timeout=2.0)
window_manager = SimpleNamespace(
event_timer_add=lambda *_args, **_kwargs: object(),
modal_handler_add=lambda *_args: None,
)
context = SimpleNamespace(
scene=SimpleNamespace(
snapform_settings=SimpleNamespace(status="", progress=0.0)
),
window_manager=window_manager,
window=object(),
)
paths = {
"log": Path("job.log"),
"status": Path("job.json"),
"result": Path("result"),
}
request = self.addon.runtime.HostJobRequest({"job_id": "job"})
try:
with patch.object(
operator,
"_prepare_command",
return_value=(request, paths),
), patch.object(
self.addon.runtime,
"submit_host_job",
side_effect=submit,
), patch.object(
self.addon.runtime,
"host_log_path",
return_value=Path("host.log"),
), patch.object(
self.addon.runtime,
"storage_root",
return_value=Path("."),
), patch.object(
self.addon.runtime,
"creation_flags",
return_value=0,
), patch.object(operator, "_redraw_ui"):
before = time.perf_counter()
result = operator.execute(context)
elapsed = time.perf_counter() - before
self.assertEqual(result, {"RUNNING_MODAL"})
self.assertLess(elapsed, 0.2)
self.assertTrue(started.wait(timeout=1.0))
finally:
release.set()
job = self.addon.operators.active_job()
if job is not None and job.start_thread is not None:
job.start_thread.join(timeout=2.0)
self.addon.operators.close_active_job()
def test_preferences_use_package_identifier(self):
self.assertEqual(
@@ -1137,11 +1235,13 @@ class AddonTest(unittest.TestCase):
Path("result"),
"BEN2_BASE",
start_frame=70,
background_depth_distance=3.0,
)
parameters = request.payload["parameters"]
self.assertEqual(parameters["video_frames"], 24)
self.assertEqual(parameters["start_frame"], 70)
self.assertEqual(parameters["background_depth_distance"], 3.0)
def test_extract_frames_request_passes_frame_range(self):
runtime = importlib.import_module("snapform.runtime")
@@ -1933,6 +2033,10 @@ class AddonTest(unittest.TestCase):
str(Path("result") / "foreground.png"),
)
self.assertFalse(calls[0][1]["snapform_delete_source"])
self.assertEqual(
calls[0][1]["snapform_depth_npz_path"],
str(Path("result") / "original-depth.npz"),
)
self.assertEqual(
calls[1],
("undo", {"message": "Convert to Depth Plane"}),

60
tests/test_image_depth.py
View File

@@ -87,6 +87,21 @@ class ImageDepthTest(unittest.TestCase):
self.assertEqual(depth.shape, (2, 3))
def test_extracts_optional_reference_depth(self):
with tempfile.TemporaryDirectory() as directory:
npz_path = Path(directory) / "results.npz"
np.savez_compressed(
npz_path,
depth=np.ones((1, 2, 2), dtype=np.float32),
snapform_reference_depth=np.float32(2.5),
)
reference_depth = self.image_depth.extract_reference_depth(
npz_path
)
self.assertEqual(reference_depth, 2.5)
def test_is_animated_image_accepts_movie_and_sequence(self):
self.assertTrue(
self.image_depth.is_animated_image(
@@ -430,6 +445,28 @@ class ImageDepthTest(unittest.TestCase):
self.assertEqual(bounds, (-1.0, 1.0, -0.5, 0.5))
def test_camera_plane_bounds_use_stored_unmasked_reference_depth(self):
source = SimpleNamespace(
data=SimpleNamespace(size=(200, 100)),
empty_image_offset=(-0.5, -0.5),
)
depth = np.full((100, 200), 9.0, dtype=np.float32)
intrinsics = np.array(
((100.0, 0.0, 100.0), (0.0, 100.0, 50.0), (0.0, 0.0, 1.0)),
dtype=np.float32,
)
bounds, reference_depth = self.image_depth.camera_plane_bounds(
source,
depth,
intrinsics,
-1.0,
reference_depth=2.0,
)
self.assertEqual(bounds, (-2.0, 2.0, -1.0, 1.0))
self.assertEqual(reference_depth, 2.0)
def test_camera_plane_bounds_preserve_square_source_aspect(self):
source = SimpleNamespace(
data=SimpleNamespace(size=(447, 447)),
@@ -460,20 +497,33 @@ class ImageDepthTest(unittest.TestCase):
def test_raw_distance_preserves_original_values(self):
depth = np.array(
((1.25, 2.5), (5.0, np.nan)),
((1.25, -2.5), (0.0, np.nan)),
dtype=np.float32,
)
mapped = self.image_depth.prepare_depth_values(depth)
self.assertEqual(float(mapped[0, 0]), 1.25)
self.assertEqual(float(mapped[0, 1]), 2.5)
self.assertEqual(float(mapped[1, 0]), 5.0)
self.assertEqual(float(mapped[0, 1]), -2.5)
self.assertEqual(float(mapped[1, 0]), 0.0)
self.assertEqual(float(mapped[1, 1]), 0.0)
def test_depth_values_are_scaled_before_creating_depth_image(self):
depth = np.array(
((0.5, -1.0), (0.0, 2.0)),
dtype=np.float32,
)
scaled = self.image_depth.scale_depth_values(depth, 4.1)
np.testing.assert_allclose(
scaled,
((2.05, -4.1), (0.0, 8.2)),
)
def test_depth_plane_uses_fixed_import_settings(self):
self.assertEqual(self.image_depth.DISPLACEMENT_STRENGTH, -1.0)
self.assertEqual(self.image_depth.DISPLACEMENT_MIDLEVEL, 1.0)
self.assertEqual(self.image_depth.DISPLACEMENT_STRENGTH, 1.0)
self.assertEqual(self.image_depth.DISPLACEMENT_MIDLEVEL, 0.0)
self.assertEqual(self.image_depth.BUMP_BASE_DISTANCE, 0.001)
def test_grid_segments_follow_aspect_ratio_from_subdivision_power(self):

110
tests/test_worker.py
View File

@@ -138,6 +138,103 @@ class WorkerTest(unittest.TestCase):
)
self.assertEqual(args.remove_background_model_dir, "ben2")
def test_background_distance_sets_fixed_background_depth(self):
prediction = SimpleNamespace(
depth=np.array(
[[
[1.0, 2.0, 3.0, 4.0],
[5.0, -9.0, -2.0, 8.0],
[9.0, -3.0, 4.0, 12.0],
[13.0, 14.0, 15.0, 16.0],
]]
),
conf=np.arange(16, dtype=np.float32).reshape(1, 4, 4),
)
foreground_mask = np.array(
[[
[True, True, True, True],
[True, True, False, True],
[True, True, True, True],
[True, True, True, True],
]],
dtype=bool,
)
original_confidence = prediction.conf.copy()
self.worker.apply_depth_background_array(
prediction,
foreground_mask,
8.0,
)
np.testing.assert_array_equal(
prediction.depth,
[[
[-8.0, -8.0, -8.0, -8.0],
[-8.0, -8.0, -8.0, -8.0],
[-8.0, -3.0, 4.0, -8.0],
[-8.0, -8.0, -8.0, -8.0],
]],
)
np.testing.assert_array_equal(
prediction.conf,
original_confidence,
)
def test_ai_depth_is_converted_for_positive_displacement(self):
prediction = SimpleNamespace(
depth=np.array([[[0.25, 1.0], [1.5, 2.0]]]),
conf=np.array([[[0.1, 0.2], [0.3, 0.4]]]),
)
self.worker.convert_ai_depth_for_displacement(prediction)
np.testing.assert_array_equal(
prediction.depth,
[[[0.75, 0.0], [-0.5, -1.0]]],
)
np.testing.assert_array_equal(
prediction.conf,
[[[0.1, 0.2], [0.3, 0.4]]],
)
def test_npz_reference_depth_preserves_prediction_arrays(self):
with tempfile.TemporaryDirectory() as directory:
npz_path = Path(directory) / "results.npz"
depth = np.array([[[1.0, 2.0]]], dtype=np.float32)
confidence = np.array([[[0.2, 0.8]]], dtype=np.float32)
np.savez_compressed(npz_path, depth=depth, conf=confidence)
self.worker.write_npz_reference_depth(npz_path, 1.5)
with np.load(npz_path, allow_pickle=False) as archive:
np.testing.assert_array_equal(archive["depth"], depth)
np.testing.assert_array_equal(archive["conf"], confidence)
self.assertEqual(
float(archive[self.worker.REFERENCE_DEPTH_KEY]),
1.5,
)
def test_original_background_uses_fixed_negative_depth(self):
luminance = np.full((3, 3), 0.5, dtype=np.float32)
alpha = np.full((3, 3), 255, dtype=np.uint8)
alpha[1, 0] = 0
depth = self.worker.compose_original_background_depth(
luminance,
alpha,
3.0,
)
np.testing.assert_allclose(
depth,
[
[-3.0, -3.0, -3.0],
[-3.0, 0.5, -3.0],
[-3.0, -3.0, -3.0],
],
)
def test_prepares_a_single_image_file(self):
with tempfile.TemporaryDirectory() as directory:
image_path = Path(directory) / "input.png"
@@ -774,7 +871,9 @@ class WorkerTest(unittest.TestCase):
def inference(self, images, **options):
inference_options.append((images, options))
return SimpleNamespace()
return SimpleNamespace(
depth=np.full((1, 2, 2), 2.0, dtype=np.float32),
)
def fake_export(prediction, export_format, export_dir, **options):
export_calls.append((prediction, export_format, export_dir, options))
@@ -862,6 +961,11 @@ 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
np.testing.assert_array_equal(
_prediction.depth,
np.full((1, 2, 2), expected_depth, dtype=np.float32),
)
if export_format == "glb":
self.assertEqual(
options,
@@ -898,7 +1002,9 @@ class WorkerTest(unittest.TestCase):
return None
def inference(self, images, **options):
return SimpleNamespace()
return SimpleNamespace(
depth=np.full((2, 2, 2), 2.0, dtype=np.float32),
)
def fake_export(prediction, export_format, export_dir, **options):
mini_directory = Path(export_dir) / "exports" / "mini_npz"