diff --git a/README.md b/README.md index 82a4514..eadfaca 100644 --- a/README.md +++ b/README.md @@ -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. 点击确认，稍等片刻就得到带位移的立体面片 diff --git a/docs/technical.md b/docs/technical.md index f41be33..12d57d0 100644 --- a/docs/technical.md +++ b/docs/technical.md @@ -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 模型时会释放旧模型；加载时如果 diff --git a/src/snapform/image_depth.py b/src/snapform/image_depth.py index 214ddb7..5f7b421 100644 --- a/src/snapform/image_depth.py +++ b/src/snapform/image_depth.py @@ -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: diff --git a/src/snapform/operators.py b/src/snapform/operators.py index f366146..b658365 100644 --- a/src/snapform/operators.py +++ b/src/snapform/operators.py @@ -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, diff --git a/src/snapform/runtime.py b/src/snapform/runtime.py index 44aab9f..89eb224 100644 --- a/src/snapform/runtime.py +++ b/src/snapform/runtime.py @@ -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) ) diff --git a/src/snapform/worker.py b/src/snapform/worker.py index 44f52af..66360ae 100644 --- a/src/snapform/worker.py +++ b/src/snapform/worker.py @@ -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( diff --git a/tests/test_addon.py b/tests/test_addon.py index b7dda58..925762b 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -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"}), diff --git a/tests/test_image_depth.py b/tests/test_image_depth.py index 1f5fc36..32c21d3 100644 --- a/tests/test_image_depth.py +++ b/tests/test_image_depth.py @@ -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): diff --git a/tests/test_worker.py b/tests/test_worker.py index 82d0c97..215ea93 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -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"