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

2
README.md
View File

@@ -30,7 +30,7 @@ DA3、BEN2 源码包和 uv 已内置在扩展里，安装时直接从本地解
2. 右键选择 `Convert to Depth Plane`
3. 在对话框中选择：
- `Depth from`：`Original` 直接用原图，`AI Generated` 让 AI 估算深度
- 需要去背景时勾选 `Remove Background`
- `Background Depth Distance` 为 `0` 时不剔除背景；大于 `0` 时把背景设为固定距离，AI 与 Original 均支持
- 其余选项保持默认即可
4. 点击确认，稍等片刻就得到带位移的立体面片

43
docs/technical.md
View File

@@ -113,26 +113,28 @@ Cyclic 沿用原 Empty 的设置。序列帧无法被 Blender 整体 Pack，因
序列帧会启动轻量抽帧任务，把帧抽出来作为 `.depth` 序列（比直接引用视频稳定）。
`Original` 是默认深度来源。选择 `AI Generated` 时，其 Model、Use Ray Pose 与 Resolution Scale
参数会集中显示在独立边框中。两种模式都有默认折叠的 Options，内部依次为 Displace via、
Subdivision、Remove Background 和 Unshaded；Delete Original 保持在折叠项外。
Remove Background 默认关闭，与 Depth from 无关；开启后需要下载 `BEN2 Base`，并从原图生成、
Pack `.foreground` 透明图片。它只替换生成面片材质的颜色图片，不修改源 Image Empty 的
Image datablock。AI Generated 仍从原图生成 `.depth`；例如 `img.png` 会得到
`img.foreground` 与 `img.depth`，不会生成 `img.foreground.depth`。Original 仍用原图作为
位移来源，不改变原有深度逻辑。
Subdivision、Background Depth Distance 和 Unshaded；Delete Original 保持在折叠项外。
Background Depth Distance 默认 `0`，表示不剔除背景。大于 `0` 时需要 `BEN2 Base`，并使用固定值
处理 alpha 背景。AI Generated 的 DA3 仍完全使用原图估算深度，推理后先把全部 depth 转换为
`1 - DA3 depth`，再只把全透明区域设为 `-Background Depth Distance`；平面尺寸仍使用转换前的深度参考值，
同时把比背景更远的前景值钳制到背景平面，并强制最外侧 1 像素为背景距离，以减少边缘墙面和毛刺；
因此背景赋值不会缩放前景。Original 则把原图前景转为灰度深度，并把全透明区域设为该距离的负值，
生成浮点 `.depth` 图片或 EXR 序列。两种模式都会使用 `.foreground` 作为材质颜色，不修改源
Image Empty 的 Image datablock，也不根据每帧前景深度动态调整背景值。
Image Empty 引用视频或序列帧时，AI Generated 会按 Empty 的 Image Frames 数值逐帧估计
深度（未设置时为完整帧数，显式设为 1 时只处理 1 帧），起始帧跟随 Empty 的 Offset，与 Remove Background 范围一致；
深度（未设置时为完整帧数，显式设为 1 时只处理 1 帧），起始帧跟随 Empty 的 Offset，与背景蒙版范围一致；
Offset 超出视频长度时同样把窗口收进视频内。材质节点与置换贴图的 Start 保持与原 Empty
一致；Original 模式（无 Remove Background）时，单帧会复制原图改名为 `.depth` 供位移使用，
视频或序列帧则抽帧保存为 job 目录根下的 `frame.0001.jpg` 等，并作为 `.depth` 序列使用。
材质颜色使用这些帧导入的 `.frame` 序列（Frames/Start/Offset 同步），开启 Remove
Background 时改用 `foreground.0001.png` 序列。AI Generated 的深度以 EXR 序列帧
（`depth.0001.exr` 起）保存在 job 目录根下；Original 模式（含 Remove Background）的位移
使用 `frame.*` 序列帧，保证逐帧对齐。
一致；Original 模式在 Distance 为 `0` 时，单帧复制原图作为 `.depth`，视频或序列帧则使用
`frame.0001.jpg` 等原图序列。Distance 大于 `0` 时改为生成包含固定负背景值的浮点 `.depth` 或
`depth.0001.exr` 序列。材质颜色在启用固定背景距离时使用 `foreground.0001.png` 序列。
AI Generated 的深度仍以 EXR 序列帧（`depth.0001.exr` 起）保存在 job 目录根下。
Displace via 可选 `Modifier`（默认）或 `Material`；
Unshaded 默认关闭；Subdivision 默认为 `5`，并默认删除原对象。
3. AI Generated 完成后会自动创建面片。面片几何先由 `mini_npz` 的相机内参与深度确定，再通过对象
Transform 匹配原 Image Empty 的可见尺寸、位置、旋转和 Offset。Original 模式直接使用源 Empty
的可见尺寸、Offset 与 Transform。开启 `Delete Original` 时，成功生成后
的可见尺寸、Offset 与 Transform。AI Generated 会把用于匹配画面尺寸的缩放预乘到 `.depth` 数值，
XY 缩放烘焙到 Mesh，Z 不参与平面匹配；最终对象 Scale 为 `(1, 1, 1)`，视觉深度与旧逻辑一致。
开启 `Delete Original` 时，成功生成后
删除源 Empty。生成面片 Object 使用源 Empty 的原名；关闭 `Delete Original` 时不处理
同名冲突，由 Blender 自动为生成面片追加编号。
4. 短边分段数为 `2^Subdivision`，长边分段数根据图像宽高比取最接近的整数。例如高:宽为
@@ -142,7 +144,7 @@ Cyclic 沿用原 Empty 的设置。序列帧无法被 Blender 整体 Pack，因
所有任务使用统一的状态栏进度显示。变化中的阶段说明显示在右侧的进度条内部，取消按钮
位于进度条右侧。状态字段支持阶段、当前操作与预计剩余秒数，界面统一组合为紧凑文案并保证
进度单调递增。Image Depth 的阶段包括 Runtime 启动、输入准备、模型加载、深度估算、缓存
进度单调递增。任务进程与 Host 请求在后台启动，按钮点击后立即进入进度状态。Image Depth 的阶段包括 Runtime 启动、输入准备、模型加载、深度估算、缓存
写出和面片构建。进度更新会强制刷新状态栏与窗口，无需移动鼠标即可实时显示。安装中断后残留的
半成品 Runtime 会在下次安装前自动清理，不会因为中断而无法再次安装。
@@ -169,15 +171,15 @@ Height = sqrt(Area / Aspect)
参考深度和尺寸。
选择 `Modifier` 时，AI Generated 的 Displace Modifier 深度贴图使用 `Extension = Extend`，
`Strength` 固定为 `-1`、`Midlevel` 固定为 `1`。Original 把原图指定给 Displace Texture，并设置
`Strength` 固定为 `1`、`Midlevel` 固定为 `0`。Original 把原图指定给 Displace Texture，并设置
`Texture Coordinates = UV`、`Midlevel = 0`；其余修改器设置保留 Blender 默认值。
视频与序列帧生成的深度贴图是 SEQUENCE Image；Displace 修改器的 Texture 与材质的 Image Texture
节点都支持序列帧，通过 image_user 的 Frames / Offset / Start 设置把第 i 帧映射到场景
`Start+i` 帧（Start 沿用 Empty 的 Start），并启用 Auto Refresh，播放时随场景帧逐帧更新。
选择 `Material` 时，不创建 Displace Modifier；深度 Image Texture 的 Color 连接到 Displacement
节点的 Height，再连接 Material Output 的 Displacement，并把材质 Surface Displacement 设置为
`Displacement Only`。AI Generated 使用 `Midlevel = 1`、`Scale = -1`；Original 使用
`Midlevel = 0`、`Scale = 1`，与对应修改器值一致。无论选择哪种呈现方式，原始彩色贴图都会连接
`Displacement Only`。AI Generated 与 Original 都使用 `Midlevel = 0`、`Scale = 1`，与对应修改器值一致。
无论选择哪种呈现方式，原始彩色贴图都会连接
Principled BSDF 的 `Base Color`，并通过
`Bump.Height → Bump.Normal → Principled.Normal` 提供表面细节。Bump 的基础 Distance 为
`0.001`，会乘以面片 Transform 的统一缩放倍率。原图 Color 还会经过 `Float Curve` 连接到
@@ -198,7 +200,7 @@ Float Curve。若 Displace via 选择 `Material`，深度贴图与 Displacement
`Convert to Depth Plane` 会按当前 Depth from 再次转换。
面片、深度 Image、Texture 与 Material 会在同一个 Blender Undo 步骤中创建；转换完成后执行
`Ctrl+Z` 会完整撤销本次转换，不会越过转换去回退更早的 Image Empty 操作。
开启 Remove Background 时，`.foreground` Image 的创建、Pack 和深度面片生成也属于
Background Depth Distance 大于 `0` 时，`.foreground` Image、Original 浮点深度、Pack 和深度面片生成也属于
同一个 Undo 步骤。
Image Empty 使用 Packed Image 时，SnapForm 会通过 Blender 写出临时 PNG 交给 DA3，
@@ -342,7 +344,8 @@ Extension 启用后只会查找已有的共享 AI Host，不会主动启动新
如果取消发生在一次不可中断的 Torch 调用中，任务会在该调用返回后的最近检查点结束。
侧边栏 `Host` 面板以 `Running (Clients: n)` 表示空闲服务，以 `Running (Busy)` 表示正在执行任务，
并显示缓存模型和活动任务。空闲且只有当前 Blender 连接时，可以清理模型缓存、重启或停止 Host；
并显示缓存模型和活动任务。Blender 完成结果导入时会立即释放当前任务对应的本地 Busy 快照，无需等待
下一次 Host 轮询。空闲且只有当前 Blender 连接时，可以清理模型缓存、重启或停止 Host；
多客户端或任务执行期间不会允许破坏共享服务。
BEN2 在首次使用后保持加载；DA3 保留最近使用的一个模型。切换 DA3 模型时会释放旧模型；加载时如果

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(

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"