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"}),