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

138
src/snapform/operators.py
View File

@@ -1,5 +1,6 @@
import shutil
import subprocess
import threading
from dataclasses import dataclass
from pathlib import Path
@@ -76,9 +77,13 @@ class ActiveJob:
displacement_via: str = "MODIFIER"
unshaded: bool = False
remove_background: bool = False
background_depth_distance: float = 0.0
temporary_input_path: Path | None = None
cancelled: bool = False
uses_host: bool = False
started: bool = False
start_error: str = ""
start_thread: threading.Thread | None = None
_active_job = None
@@ -92,12 +97,14 @@ def cancel_active_job():
if _active_job is None:
return
_active_job.cancelled = True
if not getattr(_active_job, "started", True):
return
if _active_job.uses_host:
try:
runtime.cancel_host_job(_active_job.status_path.stem)
except RuntimeError:
pass
elif _active_job.process.poll() is None:
elif _active_job.process is not None and _active_job.process.poll() is None:
_active_job.process.terminate()
@@ -190,8 +197,9 @@ class RunJob(bpy.types.Operator):
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_remove_background: bpy.props.BoolProperty(
default=False,
snapform_background_depth_distance: bpy.props.FloatProperty(
default=0.0,
min=0.0,
options={"HIDDEN", "SKIP_SAVE"},
)
@@ -205,39 +213,18 @@ class RunJob(bpy.types.Operator):
settings = context.scene.snapform_settings
try:
command, paths = self._prepare_command(context, settings)
uses_host = isinstance(command, runtime.HostJobRequest)
log_path = runtime.host_log_path() if uses_host else paths["log"]
working_directory = runtime.storage_root()
creation_flags = runtime.creation_flags()
except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
uses_host = isinstance(command, runtime.HostJobRequest)
log_file = None
log_path = paths["log"]
try:
if uses_host:
runtime.submit_host_job(command)
process = None
log_path = runtime.host_log_path()
else:
if self.job_type == "INSTALL_RUNTIME":
runtime.stop_host()
log_file = paths["log"].open("w", encoding="utf-8")
process = subprocess.Popen(
command,
stdout=log_file,
stderr=subprocess.STDOUT,
cwd=runtime.storage_root(),
creationflags=runtime.creation_flags(),
)
except (OSError, RuntimeError) as error:
if log_file is not None:
log_file.close()
self.report({"ERROR"}, f"Unable to start task: {error}")
return {"CANCELLED"}
_active_job = ActiveJob(
job = ActiveJob(
kind=self.job_type,
process=process,
log_file=log_file,
process=None,
log_file=None,
log_path=log_path,
status_path=paths["status"],
result_directory=paths["result"],
@@ -253,7 +240,8 @@ class RunJob(bpy.types.Operator):
delete_source=self.snapform_delete_source,
displacement_via=self.snapform_displacement_via,
unshaded=self.snapform_unshaded,
remove_background=self.snapform_remove_background,
remove_background=self.snapform_background_depth_distance > 0.0,
background_depth_distance=self.snapform_background_depth_distance,
temporary_input_path=(
Path(self.snapform_input_path)
if self.job_type in {
@@ -267,6 +255,7 @@ class RunJob(bpy.types.Operator):
),
uses_host=uses_host,
)
_active_job = job
settings.status = self._starting_message(settings)
settings.progress = 0.0
self._timer = context.window_manager.event_timer_add(
@@ -275,13 +264,56 @@ class RunJob(bpy.types.Operator):
)
context.window_manager.modal_handler_add(self)
self._redraw_ui(context, force=True)
job.start_thread = threading.Thread(
target=self._start_job,
args=(
job,
command,
working_directory,
creation_flags,
),
name="SnapFormJobStart",
daemon=True,
)
job.start_thread.start()
return {"RUNNING_MODAL"}
def _start_job(self, job, command, working_directory, creation_flags):
if job.cancelled:
job.started = True
return
try:
if getattr(job, "uses_host", False):
runtime.submit_host_job(command)
else:
if job.kind == "INSTALL_RUNTIME":
runtime.stop_host()
job.log_file = job.log_path.open("w", encoding="utf-8")
job.process = subprocess.Popen(
command,
stdout=job.log_file,
stderr=subprocess.STDOUT,
cwd=working_directory,
creationflags=creation_flags,
)
job.started = True
if job.cancelled:
if job.uses_host:
runtime.cancel_host_job(job.status_path.stem)
elif job.process is not None:
job.process.terminate()
except (OSError, RuntimeError) as error:
if job.log_file is not None:
job.log_file.close()
job.log_file = None
job.start_error = str(error)
def _prepare_command(self, _context, settings):
paths = runtime.new_job_paths(self.job_type)
remove_background = bool(
getattr(self, "snapform_remove_background", False)
background_depth_distance = float(
getattr(self, "snapform_background_depth_distance", 0.0)
)
remove_background = background_depth_distance > 0.0
if self.job_type == "INSTALL_RUNTIME":
self._require_online_access()
return runtime.install_runtime_command(paths["status"]), paths
@@ -338,6 +370,11 @@ class RunJob(bpy.types.Operator):
BEN2_MODEL.key,
start_frame=start_frame,
device=self._task_device(settings),
background_depth_distance=(
background_depth_distance
if self.job_type == ORIGINAL_BACKGROUND_JOB
else 0.0
),
),
paths,
)
@@ -399,6 +436,11 @@ class RunJob(bpy.types.Operator):
self.job_type == IMAGE_DEPTH_JOB
and remove_background
),
background_depth_distance=(
background_depth_distance
if self.job_type == IMAGE_DEPTH_JOB
else 0.0
),
device=self._task_device(settings),
),
paths,
@@ -455,6 +497,20 @@ class RunJob(bpy.types.Operator):
return {"CANCELLED"}
self._update_status(context, job)
if job.start_error:
settings = context.scene.snapform_settings
settings.status = "Task failed to start"
settings.progress = 0.0
self.report(
{"ERROR"},
f"Unable to start task: {job.start_error}",
)
self._remove_timer(context)
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
if not job.started:
return {"RUNNING_MODAL"}
if job.uses_host:
status = runtime.read_status(job.status_path) or {}
state = status.get("state")
@@ -532,6 +588,8 @@ class RunJob(bpy.types.Operator):
if job.cancelled:
settings.status = "Task cancelled"
settings.progress = 0.0
if getattr(job, "uses_host", False):
runtime.mark_host_job_complete(job.status_path.stem)
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
@@ -550,6 +608,8 @@ class RunJob(bpy.types.Operator):
{"ERROR"},
f"{error or fallback_error}; log: {job.log_path}",
)
if job.uses_host:
runtime.mark_host_job_complete(job.status_path.stem)
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
@@ -576,6 +636,7 @@ class RunJob(bpy.types.Operator):
settings.status = "Depth import failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
runtime.mark_host_job_complete(job.status_path.stem)
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
@@ -590,6 +651,7 @@ class RunJob(bpy.types.Operator):
settings.status = "Depth import failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
runtime.mark_host_job_complete(job.status_path.stem)
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
@@ -604,6 +666,7 @@ class RunJob(bpy.types.Operator):
settings.status = "Image update failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
runtime.mark_host_job_complete(job.status_path.stem)
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
@@ -615,6 +678,8 @@ class RunJob(bpy.types.Operator):
job.status_path.unlink(missing_ok=True)
except OSError:
pass
if getattr(job, "uses_host", False):
runtime.mark_host_job_complete(job.status_path.stem)
close_active_job()
self._redraw_ui(context, force=True)
return {"FINISHED"}
@@ -677,6 +742,11 @@ class RunJob(bpy.types.Operator):
snapform_depth_frames_dir=(
str(result_directory) if is_sequence else ""
),
snapform_depth_npz_path=(
str(result_directory / "original-depth.npz")
if job.kind == ORIGINAL_BACKGROUND_JOB
else ""
),
snapform_subdivision_level=job.subdivision_level,
snapform_delete_source=job.delete_source,
snapform_displacement_via=job.displacement_via,