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

145
src/snapform/image_depth.py
View File

@@ -7,8 +7,9 @@ import bpy
DEPTH_KEY_CANDIDATES = ("depth", "depths")
INTRINSICS_KEY_CANDIDATES = ("intrinsics", "ixts")
DISPLACEMENT_STRENGTH = -1.0
DISPLACEMENT_MIDLEVEL = 1.0
DISPLACEMENT_STRENGTH = 1.0
DISPLACEMENT_MIDLEVEL = 0.0
REFERENCE_DEPTH_KEY = "snapform_reference_depth"
ORIGINAL_DISPLACEMENT_STRENGTH = 1.0
ORIGINAL_DISPLACEMENT_MIDLEVEL = 0.0
BUMP_BASE_DISTANCE = 0.001
@@ -548,12 +549,28 @@ def apply_plane_transform(plane_object, source_object, camera_bounds):
(
(scale_x, 0.0, 0.0, offset_x),
(0.0, scale_y, 0.0, offset_y),
(0.0, 0.0, scale_z, 0.0),
(0.0, 0.0, 1.0, 0.0),
(0.0, 0.0, 0.0, 1.0),
)
)
plane_object.matrix_world = source_object.matrix_world @ fit_matrix
return abs(float(plane_object.scale.z))
return abs(float(plane_object.scale.z)) * abs(scale_z)
def apply_object_scale_to_mesh(plane_object):
"""Bake object scale into the plane while keeping location and rotation."""
from mathutils import Matrix
scale = plane_object.scale.copy()
plane_object.data.transform(
Matrix.Diagonal((scale.x, scale.y, scale.z, 1.0))
)
plane_object.data.update()
plane_object.scale = (1.0, 1.0, 1.0)
def scale_depth_values(depth, depth_scale):
return prepare_depth_values(depth) * depth_scale
def extract_depth_arrays(npz_path):
@@ -634,12 +651,31 @@ def extract_intrinsics_array(npz_path):
return intrinsics
def extract_reference_depth(npz_path):
try:
import numpy as np
except ImportError as error:
raise RuntimeError("Blender Python does not provide NumPy") from error
try:
with np.load(npz_path, allow_pickle=False) as archive:
if REFERENCE_DEPTH_KEY not in archive.files:
return None
reference_depth = float(archive[REFERENCE_DEPTH_KEY])
except OSError as error:
raise ValueError(f"Unable to read Mini NPZ: {npz_path}") from error
if not np.isfinite(reference_depth) or reference_depth <= 0.0:
raise ValueError("Mini NPZ has an invalid reference depth")
return reference_depth
def prepare_depth_values(depth):
import numpy as np
valid = np.isfinite(depth) & (depth > 0.0)
valid = np.isfinite(depth)
if not valid.any():
raise ValueError("Depth array contains no positive finite values")
raise ValueError("Depth array contains no finite values")
values = depth.astype(np.float32, copy=True)
values[~valid] = 0.0
return values
@@ -650,6 +686,7 @@ def camera_plane_bounds(
depth,
intrinsics,
displacement_strength,
reference_depth=None,
):
import numpy as np
@@ -657,10 +694,11 @@ def camera_plane_bounds(
if scale <= 1e-8:
raise ValueError("Displacement Strength must not be zero")
valid = np.isfinite(depth) & (depth > 0.0)
if not valid.any():
raise ValueError("Depth array contains no positive finite values")
reference_depth = float(np.median(depth[valid]))
if reference_depth is None:
valid = np.isfinite(depth) & (depth > 0.0)
if not valid.any():
raise ValueError("Depth array contains no positive finite values")
reference_depth = float(np.median(depth[valid]))
reference_distance = reference_depth * scale
height, width = depth.shape
@@ -978,6 +1016,7 @@ def build_image_depth_plane(
unshaded=False,
color_image=None,
depth_frames_directory=None,
depth_image=None,
):
source_name = source_object.name
frame_count = 1
@@ -1001,11 +1040,12 @@ def build_image_depth_plane(
frame_count if uses_extracted else source_frame_duration
)
texture_frame_offset = 0 if uses_extracted else source_frame_offset
depth_source_image = source_object.data
if frames_image is not None:
depth_source_image = frames_image
elif not is_animated_image(source_object.data):
depth_source_image = copy_depth_image(source_object)
depth_source_image = depth_image or source_object.data
if depth_image is None:
if frames_image is not None:
depth_source_image = frames_image
elif not is_animated_image(source_object.data):
depth_source_image = copy_depth_image(source_object)
mesh = create_plane_mesh(
source_object,
subdivision_level=subdivision_level,
@@ -1085,12 +1125,14 @@ def build_depth_plane(
depth_arrays = extract_depth_arrays(npz_path)
depth = depth_arrays[0]
intrinsics = extract_intrinsics_array(npz_path)
stored_reference_depth = extract_reference_depth(npz_path)
depth_values = prepare_depth_values(depth)
bounds, reference_depth = camera_plane_bounds(
source_object,
depth_values,
intrinsics,
DISPLACEMENT_STRENGTH,
reference_depth=stored_reference_depth,
)
depth_name = depth_name or depth_data_name(source_object)
frame_count = len(depth_arrays)
@@ -1100,18 +1142,6 @@ def build_depth_plane(
source_frame_duration,
source_use_cyclic,
) = image_user_settings(source_object)
depth_image = create_depth_image(
depth_name,
depth_values,
)
if frame_count > 1 and depth_frames_directory is not None:
bpy.data.images.remove(depth_image, do_unlink=True)
depth_image = create_depth_sequence_image(
depth_name,
[prepare_depth_values(frame) for frame in depth_arrays],
context.scene,
depth_frames_directory,
)
mesh = create_plane_mesh(source_object, bounds, subdivision_level)
plane_object = bpy.data.objects.new(
source_name,
@@ -1123,6 +1153,22 @@ def build_depth_plane(
source_object,
bounds,
)
apply_object_scale_to_mesh(plane_object)
scaled_depth_arrays = [
scale_depth_values(frame, depth_scale) for frame in depth_arrays
]
depth_image = create_depth_image(
depth_name,
scaled_depth_arrays[0],
)
if frame_count > 1 and depth_frames_directory is not None:
bpy.data.images.remove(depth_image, do_unlink=True)
depth_image = create_depth_sequence_image(
depth_name,
scaled_depth_arrays,
context.scene,
depth_frames_directory,
)
subdivision = plane_object.modifiers.new("Subdivision", "SUBSURF")
subdivision.subdivision_type = "SIMPLE"
@@ -1298,6 +1344,9 @@ class BuildOriginalDepthPlane(bpy.types.Operator):
snapform_depth_frames_dir: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_depth_npz_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_subdivision_level: bpy.props.IntProperty(
default=5,
min=0,
@@ -1342,6 +1391,23 @@ class BuildOriginalDepthPlane(bpy.types.Operator):
f"{image_base_name(source_object.data)}.frame",
self.snapform_color_frames_dir,
)
depth_image = None
if self.snapform_depth_npz_path:
depth_arrays = extract_depth_arrays(
self.snapform_depth_npz_path
)
if len(depth_arrays) == 1:
depth_image = create_depth_image(
depth_data_name(source_object),
depth_arrays[0],
)
else:
depth_image = create_depth_sequence_image(
depth_data_name(source_object),
depth_arrays,
context.scene,
self.snapform_depth_frames_dir,
)
build_image_depth_plane(
context,
source_object,
@@ -1353,6 +1419,7 @@ class BuildOriginalDepthPlane(bpy.types.Operator):
depth_frames_directory=(
self.snapform_depth_frames_dir or None
),
depth_image=depth_image,
)
except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error))
@@ -1474,10 +1541,12 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
description="Display the source color through an Emission shader",
default=False,
)
snapform_remove_background: bpy.props.BoolProperty(
name="Remove Background",
description="Remove the color image background with BEN2",
default=False,
snapform_background_depth_distance: bpy.props.FloatProperty(
name="Background Depth Distance",
description="Set excluded background to this fixed depth distance; zero disables background exclusion",
default=0.0,
min=0.0,
soft_max=10.0,
)
snapform_model: bpy.props.EnumProperty(
name="Model",
@@ -1572,7 +1641,10 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
settings_body.prop(self, "snapform_subdivision_level")
remove_background_row = settings_body.row()
remove_background_row.enabled = model_ready(BEN2_MODEL.key)
remove_background_row.prop(self, "snapform_remove_background")
remove_background_row.prop(
self,
"snapform_background_depth_distance",
)
settings_body.prop(self, "snapform_unshaded")
layout.prop(self, "snapform_delete_original")
@@ -1590,7 +1662,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
)
if (
self.snapform_depth_source == "ORIGINAL"
and not self.snapform_remove_background
and self.snapform_background_depth_distance <= 0.0
):
if video_total_frames(source_object.data) > 1:
try:
@@ -1657,6 +1729,9 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
snapform_delete_input=delete_input,
snapform_displacement_via=self.snapform_displacement_via,
snapform_unshaded=self.snapform_unshaded,
snapform_background_depth_distance=(
self.snapform_background_depth_distance
),
)
if "CANCELLED" not in result:
return {"FINISHED"}
@@ -1695,7 +1770,9 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
snapform_delete_input=delete_input,
snapform_displacement_via=self.snapform_displacement_via,
snapform_unshaded=self.snapform_unshaded,
snapform_remove_background=self.snapform_remove_background,
snapform_background_depth_distance=(
self.snapform_background_depth_distance
),
)
if "CANCELLED" in result:
if delete_input:

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,

38
src/snapform/runtime.py
View File

@@ -6,6 +6,7 @@ import shutil
import socket
import subprocess
import sys
import threading
import time
import uuid
from dataclasses import dataclass
@@ -65,6 +66,7 @@ _client_id = uuid.uuid4().hex
_host_auto_start = False
_host_snapshot = {"state": "STOPPED", "message": "Stopped"}
_local_opener = build_opener(ProxyHandler({}))
_client_label = f"Blender {getattr(bpy.app, 'version_string', '')}".strip()
def preferences():
@@ -391,11 +393,10 @@ def ensure_host(timeout=60.0):
last_launch = 0.0
starting = False
while time.monotonic() < deadline:
connection, _health = _discover_live_host(timeout=0.5)
connection, health = _discover_live_host(timeout=0.5)
if connection is not None:
_host_connection = connection
_register_host_client(connection)
health = _check_host(connection, timeout=0.5)
_host_snapshot = _format_host_snapshot(connection, health)
_redraw_host_ui()
return connection
@@ -415,7 +416,7 @@ def ensure_host(timeout=60.0):
def submit_host_job(request):
connection = ensure_host()
connection = _host_connection or ensure_host()
payload = dict(request.payload)
payload["client_id"] = _client_id
_host_http_request(
@@ -564,6 +565,24 @@ def host_busy():
return _host_snapshot.get("state") == "BUSY"
def mark_host_job_complete(job_id):
"""Release a stale local busy snapshot after its terminal status is observed."""
global _host_snapshot
active = _host_snapshot.get("active_job")
if active and active.get("job_id") != job_id:
return
if _host_snapshot.get("state") != "BUSY":
return
_host_snapshot = {
**_host_snapshot,
"busy": False,
"active_job": None,
"state": "READY",
"message": "Ready",
}
def _service_host_timer():
global _host_connection, _host_snapshot, _registered_instance_id
@@ -619,12 +638,11 @@ def _register_host_client(connection):
if _registered_instance_id == connection.instance_id:
return
label = f"Blender {getattr(bpy.app, 'version_string', '')}".strip()
_host_http_request(
connection,
"POST",
"/clients",
payload={"client_id": _client_id, "label": label},
payload={"client_id": _client_id, "label": _client_label},
timeout=2.0,
)
_registered_instance_id = connection.instance_id
@@ -641,6 +659,8 @@ def _format_host_snapshot(connection, health):
def _redraw_host_ui():
if threading.current_thread() is not threading.main_thread():
return
try:
from .status_bar import redraw_ui
@@ -701,6 +721,7 @@ def remove_background_request(
video_frames=None,
start_frame=1,
device=None,
background_depth_distance=0.0,
):
return HostJobRequest(
{
@@ -718,6 +739,9 @@ def remove_background_request(
else video_frames
),
"start_frame": start_frame,
"background_depth_distance": float(
background_depth_distance
),
},
}
)
@@ -761,6 +785,7 @@ def reconstruct_request(
process_res_method=None,
use_ray_pose=None,
remove_background=False,
background_depth_distance=0.0,
start_frame=1,
device=None,
):
@@ -798,8 +823,9 @@ def reconstruct_request(
"use_ray_pose": 1 if snapform_selected_ray_pose else 0,
"ref_view_strategy": settings.ref_view_strategy,
"remove_background_model_dir": "",
"background_depth_distance": float(background_depth_distance),
}
if remove_background:
if remove_background or background_depth_distance > 0.0:
parameters["remove_background_model_dir"] = str(
model_directory(BEN2_MODEL.key)
)

175
src/snapform/worker.py
View File

@@ -29,6 +29,7 @@ HF_ENDPOINT = "https://huggingface.co"
MODELSCOPE_ENDPOINT = "https://modelscope.cn"
MODELSCOPE_REVISION = "master"
MODELSCOPE_TIMEOUT = 60.0
REFERENCE_DEPTH_KEY = "snapform_reference_depth"
class TaskCancelled(RuntimeError):
@@ -40,6 +41,144 @@ def check_cancelled(cancel_check):
raise TaskCancelled("Task cancelled")
def apply_foreground_depth_distance(
prediction,
foreground_paths,
background_depth_distance,
):
"""Set fully transparent pixels to a fixed negative displacement."""
import numpy as np
from PIL import Image
if len(foreground_paths) != len(prediction.depth):
raise RuntimeError("Foreground mask count does not match depth frames")
masks = []
height, width = prediction.depth.shape[-2:]
for foreground_path in foreground_paths:
with Image.open(foreground_path) as source:
alpha = source.convert("RGBA").getchannel("A")
alpha = alpha.resize((width, height), resample=Image.Resampling.NEAREST)
masks.append(np.asarray(alpha) > 0)
apply_depth_background_array(
prediction,
np.stack(masks),
background_depth_distance,
)
def apply_depth_background_array(
prediction,
foreground_mask,
background_depth_distance,
):
import numpy as np
depth = np.asarray(prediction.depth)
masked_depth = depth.copy()
background_depth = -background_depth_distance
for frame_index in range(len(masked_depth)):
np.maximum(
masked_depth[frame_index],
background_depth,
out=masked_depth[frame_index],
)
masked_depth[frame_index][~foreground_mask[frame_index]] = (
background_depth
)
set_depth_border(masked_depth[frame_index], background_depth)
prediction.depth = masked_depth
def set_depth_border(depth, background_depth):
"""Set the outermost pixel ring to the fixed background depth."""
if depth.ndim != 2 or not depth.size:
raise ValueError("Depth frame must be a non-empty 2D array")
depth[0, :] = background_depth
depth[-1, :] = background_depth
depth[:, 0] = background_depth
depth[:, -1] = background_depth
def convert_ai_depth_for_displacement(prediction):
"""Convert DA3 distance into Blender's positive displacement direction."""
import numpy as np
prediction.depth = 1.0 - np.asarray(prediction.depth)
def depth_reference_value(depth):
"""Return the unmasked depth used to size the imported camera plane."""
import numpy as np
values = np.asarray(depth)
valid = np.isfinite(values) & (values > 0.0)
if not valid.any():
raise ValueError("Depth array contains no positive finite values")
return float(np.median(values[valid]))
def write_npz_reference_depth(npz_path, reference_depth):
"""Add SnapForm metadata without changing exported prediction arrays."""
import numpy as np
npz_path = Path(npz_path)
temporary_path = npz_path.with_name(f"{npz_path.stem}.tmp.npz")
with np.load(npz_path, allow_pickle=False) as archive:
values = {key: archive[key] for key in archive.files}
values[REFERENCE_DEPTH_KEY] = np.asarray(reference_depth, dtype=np.float32)
np.savez_compressed(temporary_path, **values)
temporary_path.replace(npz_path)
def write_original_background_depth(
image_paths,
foreground_paths,
background_depth_distance,
output_directory,
):
"""Write original-image luminance with a fixed negative background depth."""
import numpy as np
from PIL import Image
if len(image_paths) != len(foreground_paths):
raise RuntimeError("Foreground mask count does not match source frames")
depth_frames = []
for image_path, foreground_path in zip(image_paths, foreground_paths):
with Image.open(image_path) as source:
luminance = np.asarray(source.convert("L"), dtype=np.float32) / 255.0
with Image.open(foreground_path) as foreground:
alpha = np.asarray(
foreground.convert("RGBA").getchannel("A"),
dtype=np.uint8,
)
if luminance.shape != alpha.shape:
raise RuntimeError("Foreground mask size does not match source frame")
depth_frames.append(
compose_original_background_depth(
luminance,
alpha,
background_depth_distance,
)
)
np.savez_compressed(
Path(output_directory) / "original-depth.npz",
depth=np.stack(depth_frames),
)
def compose_original_background_depth(
luminance,
alpha,
background_depth_distance,
):
import numpy as np
background_depth = -background_depth_distance
depth = np.where(alpha > 0, luminance, background_depth)
set_depth_border(depth, background_depth)
return depth
def write_status(
path,
progress,
@@ -757,6 +896,21 @@ def remove_background(args, model_cache=None, cancel_check=None):
cancel_check=cancel_check,
)
copy_frames(images, output_directory)
background_depth_distance = float(
getattr(args, "background_depth_distance", 0.0)
)
if background_depth_distance > 0.0:
foreground_paths = (
[output_directory / "foreground.png"]
if len(images) == 1
else sorted(output_directory.glob("foreground.*.png"))
)
write_original_background_depth(
images,
foreground_paths,
background_depth_distance,
output_directory,
)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
@@ -878,6 +1032,10 @@ def reconstruct(args, model_cache=None, cancel_check=None):
"remove_background_model_dir",
"",
)
background_depth_distance = float(
getattr(args, "background_depth_distance", 0.0)
)
foreground_images = []
if remove_background_model_dir:
write_status(args.status, 0.12, "Removing image background")
if len(images) == 1:
@@ -897,6 +1055,7 @@ def reconstruct(args, model_cache=None, cancel_check=None):
)
if not (output_directory / "foreground.png").is_file():
raise RuntimeError("BEN2 did not write a foreground image")
foreground_images = [output_directory / "foreground.png"]
else:
load_ms, inference_ms = infer_foreground_frames(
images,
@@ -910,6 +1069,10 @@ def reconstruct(args, model_cache=None, cancel_check=None):
model_cache=model_cache,
cancel_check=cancel_check,
)
foreground_images = sorted(
output_directory.glob("foreground.*.png")
)
if remove_background_model_dir:
print(
f"BEN2 inference {inference_ms:.0f} ms; "
f"model load {load_ms:.0f} ms"
@@ -950,6 +1113,16 @@ def reconstruct(args, model_cache=None, cancel_check=None):
ref_view_strategy=args.ref_view_strategy,
)
check_cancelled(cancel_check)
reference_depth = None
if args.export_format == "mini_npz":
reference_depth = depth_reference_value(prediction.depth[0])
convert_ai_depth_for_displacement(prediction)
if background_depth_distance > 0.0:
apply_foreground_depth_distance(
prediction,
foreground_images,
background_depth_distance,
)
write_status(args.status, 0.88, f"Exporting {args.export_format}")
export(
prediction,
@@ -981,6 +1154,8 @@ def reconstruct(args, model_cache=None, cancel_check=None):
directory.rmdir()
except OSError:
pass
if reference_depth is not None:
write_npz_reference_depth(result_path, reference_depth)
check_cancelled(cancel_check)
copy_frames(images, output_directory)
write_status(