From d694980da7b6a650c224d0b65fa17e34de400e86 Mon Sep 17 00:00:00 2001 From: nan Date: Mon, 3 Aug 2026 20:18:00 +0800 Subject: [PATCH] feat: Add BEN2 background removal functionality and integrate into SnapForm - Introduced BEN2 model for background removal in the panel and properties. - Updated runtime to handle background removal commands and model downloads. - Implemented background removal logic in worker, including image processing and saving results. - Enhanced tests to cover new background removal features and ensure proper integration. - Updated version to 0.9.28 to reflect new features and improvements. --- README.md | 53 +++- pyproject.toml | 2 +- src/snapform/__init__.py | 6 + src/snapform/blender_manifest.toml | 2 +- src/snapform/bootstrap_runtime.py | 21 +- src/snapform/image_depth.py | 375 +++++++++++++++++++--- src/snapform/models.py | 19 ++ src/snapform/operators.py | 178 ++++++++++- src/snapform/panel.py | 15 +- src/snapform/properties.py | 5 + src/snapform/runtime.py | 49 ++- src/snapform/worker.py | 112 ++++++- tests/test_addon.py | 494 ++++++++++++++++++++++++++++- tests/test_bootstrap_runtime.py | 17 + tests/test_image_depth.py | 87 ++++- tests/test_models.py | 11 + tests/test_pack.py | 2 +- tests/test_worker.py | 157 +++++++++ uv.lock | 2 +- 19 files changed, 1523 insertions(+), 84 deletions(-) diff --git a/README.md b/README.md index ecd68b2..007867b 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,8 @@ `SnapForm` 是一套从图片快速构建空间结构的 Blender Extension。目前以 [Depth Anything 3](https://github.com/ByteDance-Seed/Depth-Anything-3) 作为首个 AI 能力，在独立 Runtime 中运行 DA3，可将官方导出结果保存在独立 job 目录， -也可把 Image Empty 转换成带几何与材质置换的立体面片。 +也可把 Image Empty 转换成带几何与材质置换的立体面片。Debug 面板还可运行 +[BEN2](https://github.com/PramaLLC/BEN2) 背景剔除功能测试。 ## 当前功能 @@ -11,11 +12,12 @@ - 在独立 Python Runtime 中运行 Torch 与 DA3，避免污染 Blender Python 进程 - Runtime、模型和任务结果默认保存在 `~/.snapform`，可由多个 Blender 版本共享 - 支持 DA3 Small、Base、Large 1.1、Giant 1.1 和 Nested Giant Large 1.1 +- 支持 BEN2 Base 单图背景剔除测试，输出透明 PNG 与推理耗时 - 支持单张 JPG、JPEG、PNG、WEBP 图片或常见视频文件输入 - 视频均匀抽帧，可设置参与重建的最大帧数 - `Setup` 集中管理 Runtime，并逐项列出和下载所有可用模型 - Runtime 已就绪时不额外占用面板空间；模型列表每个模型只占一行 -- `Debug` 是独立且默认折叠的 DA3 测试面板 +- `Debug` 是独立且默认折叠的 AI 功能测试面板 - 使用一个导出类型下拉菜单和一个 `Export` 按钮 - 只显示当前格式的专属参数；GLB 与 PLY 设置不会同时出现 - Debug 的模型选项按格式过滤；PLY 只提供两个带 Gaussian head 的 Giant 模型 @@ -23,10 +25,11 @@ `(Missing)`，已下载模型保持原名 - 分别调用 DA3 的 `glb`、`gs_ply`、`depth_vis` 和 `mini_npz` exporter - Image Empty 右键菜单可直接运行 DA3，根据 `mini_npz` 深度生成面片 +- Image Empty 右键菜单可直接运行 BEN2，保留原图并换用新建的透明背景图片 - 使用浮点深度贴图、Simple Subdivision、Displace Modifier 和原图 Bump 构建立体面片 - 转换对话框可选择推理模式、处理分辨率和面片网格细分 -- Debug 导出完成后只打开对应 job 文件夹，不向当前 Blender 文件添加任何数据 +- Debug 的 DA3 导出或 BEN2 背景剔除完成后只打开对应 job 文件夹，不向当前 Blender 文件添加任何数据 - 完整保留每次任务生成的所有官方文件 - 安装、模型下载和推理均在后台子进程运行，状态栏实时显示阶段、百分比进度条与取消按钮 - 后台任务运行时，Install、Download 和 Export 按钮保持原位置并变灰，任务完成后恢复 @@ -36,11 +39,27 @@ 1. 使用 `uv run python pack.py` 构建 Extension ZIP。 2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。 3. 在 Extension Preferences 中确认 `uv` 路径和 Torch Backend。 -4. 点击 `Install Runtime`。首次安装会下载 Python、PyTorch 和 DA3 依赖。 +4. 点击 `Install Runtime`。首次安装会下载 Python、PyTorch、DA3 和 BEN2 依赖。 5. 在 3D View 按 `N`，打开 `SnapForm` 页签。 6. 在 `Setup` 中找到需要的模型并点击该行的 `Download`。 7. 展开默认收起的 `Debug`，选择 `Export Format` 和本次执行使用的模型。 8. 选择一张图片或一个视频文件。视频可设置均匀抽取的帧数，然后点击 `Export`。 +9. 测试背景剔除时，在 `Setup` 下载 `BEN2 Base`，然后在 `Debug` 的 + `BEN2 Background Removal` 区域选择单张图片并点击 `Remove Background`。 + 完成后会打开 job 文件夹，其中的 `foreground.png` 带透明 Alpha；状态栏显示 + BEN2 推理、模型加载与任务总耗时。当前测试固定关闭较慢的 foreground refinement。 + +## Image Empty 去背景 + +1. 在 `Setup` 下载 `BEN2 Base`。 +2. 在 3D View 选中 Image Empty，右键选择 `Remove Background`。 +3. 完成后以 Fake User 保留原 Image datablock，新建并 Pack 一张透明 PNG 到当前 `.blend`，然后只把 + 当前 Empty 的 `data` 换成新图；Empty 的 Transform、显示尺寸、Offset 及其他设置保持不变。 +4. 新 Image 使用原图去掉格式后再追加 `.rb` 的名称。例如 `img.png` 生成 `img.rb`， + `img.png.001` 生成 `img.001.rb`。已有同名 datablock 时由 Blender 自动追加编号。 + +原图已 Pack 时会先写出临时 PNG 交给 BEN2，任务结束后删除临时输入；原 Image 仍保持 Pack +状态和原始内容。BEN2 的 `foreground.png` 也会保留在对应 job 目录中。 ## Image Empty 转深度面片 @@ -52,11 +71,17 @@ 转换导出 `mini_npz`，使用前需要在 `Setup` 下载所选模型，并固定使用 `process_res_method="upper_bound_resize"`；设备和参考视图策略 继续使用场景设置。 - - `Original`：直接把原图作为深度贴图，不显示估算参数，也不启动 Runtime 或要求源图片存在于磁盘。 - `Original` 是默认深度来源。选择 `AI Generated` 时，其 Model、Use Ray Pose 和 Resolution Scale - 参数会集中显示在独立边框中。两种模式始终显示 Displace via、Subdivision 与 Delete Original。 + - `Original`：直接把原图作为深度贴图，不显示估算参数。Remove Background 关闭时不启动 + Runtime，也不要求源图片存在于磁盘。 + `Original` 是默认深度来源。选择 `AI Generated` 时，其 Model、Use Ray Pose 与 Resolution Scale + 参数会集中显示在独立边框中。两种模式始终显示 Remove Background、Displace via、Unshaded、 + Subdivision 与 Delete Original。 + Remove Background 默认关闭，与 Depth from 无关；开启后需要下载 `BEN2 Base`，并从原图生成、 + Pack `.rb` 透明图片。它只替换生成面片材质的颜色图片，不修改源 Image Empty 的 Image datablock。 + AI Generated 仍从原图生成 `.depth`；例如 `img.png` 会得到 `img.rb` 与 `img.depth`，不会生成 + `img.rb.depth`。Original 仍用原图作为位移来源，不改变原有深度逻辑。 Displace via 可选 `Modifier`（默认）或 `Material`； - Subdivision 默认为 `5`，并默认删除原对象。 + Unshaded 默认关闭；Subdivision 默认为 `5`，并默认删除原对象。 3. AI Generated 完成后会自动创建面片。面片几何先由 `mini_npz` 的相机内参与深度确定，再通过对象 Transform 匹配原 Image Empty 的可见尺寸、位置、旋转和 Offset。Original 模式直接使用源 Empty 的可见尺寸、Offset 与 Transform。开启 `Delete Original` 时，成功生成后 @@ -105,6 +130,9 @@ Principled BSDF 的 `Base Color`，并通过 `0.001`，会乘以面片 Transform 的统一缩放倍率。原图 Color 还会经过 `Float Curve` 连接到 Principled BSDF 的 `Roughness`，曲线两端输出分别为 `0.5` 和 `0.1`；材质 IOR 保持 Principled BSDF 默认值 `1.5`。 +开启 `Unshaded` 后，Surface 改为原图 Color 直接连接 Emission，再连接 Material Output；不创建 +Principled BSDF、Bump 或 Float Curve。若 Displace via 选择 `Material`，深度贴图与 Displacement +节点仍会保留并连接 Material Output 的 Displacement。 这里的“原始距离”是 DA3 模型输出的尺度。普通 Small、Base、Large 和 Giant 模型的深度不应 被解释为现实世界的米；需要现实尺度时应选择带 metric depth 的 @@ -116,6 +144,8 @@ Principled BSDF 默认值 `1.5`。 `Convert to Depth Plane` 会按当前 Depth from 再次转换。 面片、深度 Image、Texture 与 Material 会在同一个 Blender Undo 步骤中创建；转换完成后执行 `Ctrl+Z` 会完整撤销本次转换，不会越过转换去回退更早的 Image Empty 操作。 +开启 Remove Background 时，`.rb` Image 的创建、Pack 和深度面片生成也属于 +同一个 Undo 步骤。 Image Empty 使用 Packed Image 时，SnapForm 会通过 Blender 写出临时 PNG 交给 DA3， 不会 Unpack 或修改源 Image；任务结束后自动删除临时文件。普通图片仍直接读取源文件。DA3 的原始 @@ -129,7 +159,7 @@ Image Empty 使用 Packed Image 时，SnapForm 会通过 Blender 写出临时 PN `DA3 Nested Giant Large 1.1`。Small、Base 和 Large 没有 Gaussian head，无法由 DA3 生成 `gs_ply`。 -每次任务会在 `~/.snapform/jobs/<任务目录>/` 中保留 DA3 官方产物： +每次任务会在 `~/.snapform/jobs/<任务目录>/` 中保留输出产物： ```text / @@ -143,6 +173,9 @@ DA3 生成 `gs_ply`。 / └── exports/mini_npz/results.npz + +/ +└── foreground.png ``` 建议从 `DA3 Base` 开始。它使用 Apache 2.0 模型权重；`DA3 Large 1.1` 权重采用 @@ -185,6 +218,7 @@ Extension 本身不包含 Torch。`Install Runtime` 使用 `uv` 创建独立 Pyt - TorchVision 0.24.1 - xFormers 0.0.33.post2（与 PyTorch 2.9.1 对应） - 固定到已验证 commit 的 Depth Anything 3 +- 固定到已验证 commit 的 BEN2 - 固定到 DA3 官方推荐 commit 的 `gsplat`；仅安装 PLY 导出所需的纯 Python 部分， 不编译本机 CUDA rasterizer - `addict` 2.4.0（补充 DA3 上游包元数据暂未声明的运行依赖） @@ -236,5 +270,6 @@ uv run python pack.py ## 上游与许可证 - Depth Anything 3 代码采用 Apache 2.0。 +- BEN2 代码与模型权重采用 MIT License。 - 本 Extension 代码采用 GPL-3.0-or-later。 - 模型权重许可证由所选 Hugging Face 模型仓库决定，Extension 下载前会在界面显示。 diff --git a/pyproject.toml b/pyproject.toml index 59ee571..30d93a6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "snapform" -version = "0.9.27" +version = "0.9.28" description = "Fast image-guided structure building tools for Blender" readme = "README.md" requires-python = ">=3.11" diff --git a/src/snapform/__init__.py b/src/snapform/__init__.py index 6493c21..d80124f 100644 --- a/src/snapform/__init__.py +++ b/src/snapform/__init__.py @@ -1,8 +1,11 @@ import bpy from .image_depth import ( + ApplyRemovedBackground, BuildDepthPlane, + BuildOriginalDepthPlane, ConvertImageToDepthPlane, + RemoveImageBackground, draw_image_context_menu, ) from .operators import ( @@ -22,7 +25,10 @@ CLASSES = ( SnapFormSettings, RunJob, CancelJob, + ApplyRemovedBackground, + RemoveImageBackground, BuildDepthPlane, + BuildOriginalDepthPlane, ConvertImageToDepthPlane, RuntimeModelsPanel, DebugPanel, diff --git a/src/snapform/blender_manifest.toml b/src/snapform/blender_manifest.toml index 00535dd..633be26 100644 --- a/src/snapform/blender_manifest.toml +++ b/src/snapform/blender_manifest.toml @@ -1,7 +1,7 @@ schema_version = "1.0.0" id = "snapform" -version = "0.9.27" +version = "0.9.28" name = "SnapForm" tagline = "Build spatial structures from images, fast" maintainer = "OmooLab" diff --git a/src/snapform/bootstrap_runtime.py b/src/snapform/bootstrap_runtime.py index fe4c0ed..1e91425 100644 --- a/src/snapform/bootstrap_runtime.py +++ b/src/snapform/bootstrap_runtime.py @@ -11,7 +11,7 @@ TORCHVISION_VERSION = "0.24.1" XFORMERS_VERSION = "0.0.33.post2" ADDICT_VERSION = "2.4.0" GSPLAT_COMMIT = "0b4dddf04cb687367602c01196913cde6a743d70" -RUNTIME_SCHEMA = 2 +RUNTIME_SCHEMA = 3 TORCH_INDEXES = { "cu128": "https://download.pytorch.org/whl/cu128", "cpu": "https://download.pytorch.org/whl/cpu", @@ -113,6 +113,22 @@ def install(args): ] ) + write_status(args.status, 0.68, "Installing BEN2") + ben2_requirement = ( + "ben2 @ git+https://github.com/PramaLLC/BEN2.git" + f"@{args.ben2_commit}" + ) + run( + [ + args.uv, + "pip", + "install", + "--python", + python, + ben2_requirement, + ] + ) + write_status(args.status, 0.78, "Installing gsplat build tools") run( [ @@ -154,6 +170,7 @@ def install(args): ( "import torch; " "from depth_anything_3.api import DepthAnything3; " + "from ben2 import BEN_Base; " "import gsplat; " "print(torch.__version__, torch.cuda.is_available())" ), @@ -170,6 +187,7 @@ def install(args): "xformers": XFORMERS_VERSION, "backend": args.backend, "da3_commit": args.da3_commit, + "ben2_commit": args.ben2_commit, "gsplat_commit": GSPLAT_COMMIT, "gsplat_cuda": False, }, @@ -186,6 +204,7 @@ def parse_args(): parser.add_argument("--runtime-dir", required=True) parser.add_argument("--backend", choices=tuple(TORCH_INDEXES), required=True) parser.add_argument("--da3-commit", required=True) + parser.add_argument("--ben2-commit", required=True) parser.add_argument("--status", required=True) return parser.parse_args() diff --git a/src/snapform/image_depth.py b/src/snapform/image_depth.py index 10729ad..97494d1 100644 --- a/src/snapform/image_depth.py +++ b/src/snapform/image_depth.py @@ -80,6 +80,10 @@ def depth_data_name(source_object): return f"{image_base_name(source_object.data)}.depth" +def removed_background_data_name(source_object): + return f"{image_base_name(source_object.data)}.rb" + + def scaled_process_resolution(image, scale): width, height = image.size if width <= 0 or height <= 0: @@ -108,6 +112,49 @@ def save_packed_image(image, scene): return output_path +def prepare_image_input(image, scene): + source_path = image_path(image) + if image_is_packed(image): + return save_packed_image(image, scene), True + if not source_path.is_file(): + raise RuntimeError( + "The image empty must reference an existing image file" + ) + return source_path, False + + +def create_removed_background_image( + source_object, + output_path, + assign_to_source=True, +): + output_path = Path(output_path) + if not output_path.is_file(): + raise RuntimeError( + f"Background removal result does not exist: {output_path}" + ) + try: + result_image = bpy.data.images.load( + str(output_path), + check_existing=False, + ) + except RuntimeError as error: + raise RuntimeError("Unable to load the background result") from error + source_image = source_object.data + original_fake_user = source_image.use_fake_user + try: + result_image.name = removed_background_data_name(source_object) + result_image.pack() + if assign_to_source: + source_image.use_fake_user = True + source_object.data = result_image + except (OSError, RuntimeError) as error: + source_image.use_fake_user = original_fake_user + bpy.data.images.remove(result_image) + raise RuntimeError("Unable to pack the background result") from error + return result_image + + def plane_bounds(source_object): image_width, image_height = source_object.data.size if image_width <= 0 or image_height <= 0: @@ -422,6 +469,7 @@ def create_material( displacement_image=None, displacement_midlevel=0.0, displacement_scale=1.0, + unshaded=False, ): material = bpy.data.materials.new(name) material.use_nodes = True @@ -430,42 +478,74 @@ def create_material( nodes = material.node_tree.nodes nodes.clear() - output = nodes.new("ShaderNodeOutputMaterial") - shader = nodes.new("ShaderNodeBsdfPrincipled") color_texture = nodes.new("ShaderNodeTexImage") - bump = nodes.new("ShaderNodeBump") - roughness_curve = nodes.new("ShaderNodeFloatCurve") + output = nodes.new("ShaderNodeOutputMaterial") color_texture.image = color_image color_texture.interpolation = "Linear" color_texture.extension = "EXTEND" - bump.inputs["Distance"].default_value = BUMP_BASE_DISTANCE * depth_scale - roughness_curve.mapping.initialize() - curve_points = roughness_curve.mapping.curves[0].points - curve_points[0].location = (0.0, 0.5) - curve_points[-1].location = (1.0, 0.1) - roughness_curve.mapping.update() - color_texture.location = (-640.0, 180.0) - roughness_curve.location = (-300.0, 320.0) - bump.location = (-300.0, -80.0) - shader.location = (80.0, 180.0) output.location = (460.0, 180.0) links = material.node_tree.links - links.new(color_texture.outputs["Color"], shader.inputs["Base Color"]) - links.new(color_texture.outputs["Alpha"], shader.inputs["Alpha"]) - links.new(color_texture.outputs["Color"], bump.inputs["Height"]) - links.new(bump.outputs["Normal"], shader.inputs["Normal"]) - links.new( - color_texture.outputs["Color"], - roughness_curve.inputs["Value"], - ) - links.new( - roughness_curve.outputs["Value"], - shader.inputs["Roughness"], - ) - links.new(shader.outputs["BSDF"], output.inputs["Surface"]) + if unshaded: + emission = nodes.new("ShaderNodeEmission") + transparent = nodes.new("ShaderNodeBsdfTransparent") + mix_shader = nodes.new("ShaderNodeMixShader") + emission.location = (-100.0, 260.0) + transparent.location = (-100.0, 20.0) + mix_shader.location = (200.0, 180.0) + links.new( + color_texture.outputs["Color"], + emission.inputs["Color"], + ) + links.new( + color_texture.outputs["Alpha"], + mix_shader.inputs["Fac"], + ) + links.new( + transparent.outputs["BSDF"], + mix_shader.inputs[1], + ) + links.new( + emission.outputs["Emission"], + mix_shader.inputs[2], + ) + links.new( + mix_shader.outputs["Shader"], + output.inputs["Surface"], + ) + else: + shader = nodes.new("ShaderNodeBsdfPrincipled") + bump = nodes.new("ShaderNodeBump") + roughness_curve = nodes.new("ShaderNodeFloatCurve") + bump.inputs["Distance"].default_value = BUMP_BASE_DISTANCE * depth_scale + roughness_curve.mapping.initialize() + curve_points = roughness_curve.mapping.curves[0].points + curve_points[0].location = (0.0, 0.5) + curve_points[-1].location = (1.0, 0.1) + roughness_curve.mapping.update() + + roughness_curve.location = (-300.0, 320.0) + bump.location = (-300.0, -80.0) + shader.location = (80.0, 180.0) + + links.new( + color_texture.outputs["Color"], + shader.inputs["Base Color"], + ) + links.new(color_texture.outputs["Alpha"], shader.inputs["Alpha"]) + links.new(color_texture.outputs["Color"], bump.inputs["Height"]) + links.new(bump.outputs["Normal"], shader.inputs["Normal"]) + links.new( + color_texture.outputs["Color"], + roughness_curve.inputs["Value"], + ) + links.new( + roughness_curve.outputs["Value"], + shader.inputs["Roughness"], + ) + links.new(shader.outputs["BSDF"], output.inputs["Surface"]) if displacement_image is not None: depth_texture = nodes.new("ShaderNodeTexImage") @@ -525,6 +605,8 @@ def build_image_depth_plane( subdivision_level, delete_source, displacement_via="MODIFIER", + unshaded=False, + color_image=None, ): source_name = source_object.name mesh = create_plane_mesh( @@ -554,13 +636,14 @@ def build_image_depth_plane( depth_scale = abs(float(plane_object.scale.z)) material = create_material( source_name, - source_object.data, + color_image if color_image is not None else source_object.data, depth_scale, displacement_image=( source_object.data if displacement_via == "MATERIAL" else None ), displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL, displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH, + unshaded=unshaded, ) mesh.materials.append(material) return finish_depth_plane( @@ -578,6 +661,9 @@ def build_depth_plane( subdivision_level, delete_source, displacement_via="MODIFIER", + unshaded=False, + depth_name="", + color_image=None, ): source_name = source_object.name depth = extract_depth_array(npz_path) @@ -589,7 +675,7 @@ def build_depth_plane( intrinsics, DISPLACEMENT_STRENGTH, ) - depth_name = depth_data_name(source_object) + depth_name = depth_name or depth_data_name(source_object) depth_image = create_depth_image( depth_name, depth_values, @@ -627,13 +713,14 @@ def build_depth_plane( material = create_material( source_name, - source_object.data, + color_image if color_image is not None else source_object.data, depth_scale, displacement_image=( depth_image if displacement_via == "MATERIAL" else None ), displacement_midlevel=DISPLACEMENT_MIDLEVEL, displacement_scale=DISPLACEMENT_STRENGTH, + unshaded=unshaded, ) mesh.materials.append(material) plane_object["snapform_reference_depth"] = reference_depth @@ -660,6 +747,12 @@ class BuildDepthPlane(bpy.types.Operator): snapform_npz_path: bpy.props.StringProperty( options={"HIDDEN", "SKIP_SAVE"}, ) + snapform_depth_data_name: bpy.props.StringProperty( + options={"HIDDEN", "SKIP_SAVE"}, + ) + snapform_background_path: bpy.props.StringProperty( + options={"HIDDEN", "SKIP_SAVE"}, + ) snapform_subdivision_level: bpy.props.IntProperty( default=5, min=0, @@ -678,6 +771,10 @@ class BuildDepthPlane(bpy.types.Operator): default="MODIFIER", options={"HIDDEN", "SKIP_SAVE"}, ) + snapform_unshaded: bpy.props.BoolProperty( + default=False, + options={"HIDDEN", "SKIP_SAVE"}, + ) def execute(self, context): source_object = bpy.data.objects.get(self.snapform_source_object_name) @@ -697,13 +794,23 @@ class BuildDepthPlane(bpy.types.Operator): return {"CANCELLED"} try: + color_image = None + if self.snapform_background_path: + color_image = create_removed_background_image( + source_object, + self.snapform_background_path, + assign_to_source=False, + ) build_depth_plane( context, source_object, npz_path, self.snapform_subdivision_level, self.snapform_delete_source, - self.snapform_displacement_via, + displacement_via=self.snapform_displacement_via, + unshaded=self.snapform_unshaded, + depth_name=self.snapform_depth_data_name, + color_image=color_image, ) except (OSError, RuntimeError, ValueError) as error: self.report({"ERROR"}, str(error)) @@ -711,6 +818,138 @@ class BuildDepthPlane(bpy.types.Operator): return {"FINISHED"} +class BuildOriginalDepthPlane(bpy.types.Operator): + bl_idname = "snapform.build_original_depth_plane" + bl_label = "Build Original Depth Plane" + bl_options = {"INTERNAL"} + + snapform_source_object_name: bpy.props.StringProperty( + options={"HIDDEN", "SKIP_SAVE"}, + ) + snapform_background_path: bpy.props.StringProperty( + options={"HIDDEN", "SKIP_SAVE"}, + ) + snapform_subdivision_level: bpy.props.IntProperty( + default=5, + min=0, + max=10, + options={"HIDDEN", "SKIP_SAVE"}, + ) + snapform_delete_source: bpy.props.BoolProperty( + default=True, + options={"HIDDEN", "SKIP_SAVE"}, + ) + snapform_displacement_via: bpy.props.EnumProperty( + items=( + ("MODIFIER", "Modifier", "Use a Displace modifier"), + ("MATERIAL", "Material", "Use material displacement"), + ), + default="MODIFIER", + options={"HIDDEN", "SKIP_SAVE"}, + ) + snapform_unshaded: bpy.props.BoolProperty( + default=False, + options={"HIDDEN", "SKIP_SAVE"}, + ) + + def execute(self, context): + source_object = bpy.data.objects.get(self.snapform_source_object_name) + if not is_image_empty(source_object): + self.report( + {"ERROR"}, + "The source image empty is no longer available", + ) + return {"CANCELLED"} + try: + color_image = create_removed_background_image( + source_object, + self.snapform_background_path, + assign_to_source=False, + ) + build_image_depth_plane( + context, + source_object, + self.snapform_subdivision_level, + self.snapform_delete_source, + displacement_via=self.snapform_displacement_via, + unshaded=self.snapform_unshaded, + color_image=color_image, + ) + except (OSError, RuntimeError, ValueError) as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + return {"FINISHED"} + + +class ApplyRemovedBackground(bpy.types.Operator): + bl_idname = "snapform.apply_removed_background" + bl_label = "Apply Removed Background" + bl_options = {"INTERNAL"} + + snapform_source_object_name: bpy.props.StringProperty( + options={"HIDDEN", "SKIP_SAVE"}, + ) + snapform_output_path: bpy.props.StringProperty( + options={"HIDDEN", "SKIP_SAVE"}, + ) + + def execute(self, _context): + source_object = bpy.data.objects.get(self.snapform_source_object_name) + if not is_image_empty(source_object): + self.report( + {"ERROR"}, + "The source image empty is no longer available", + ) + return {"CANCELLED"} + try: + create_removed_background_image( + source_object, + self.snapform_output_path, + ) + except RuntimeError as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + return {"FINISHED"} + + +class RemoveImageBackground(bpy.types.Operator): + bl_idname = "snapform.remove_image_background" + bl_label = "Remove Background" + bl_description = "Remove the background from this Image Empty with BEN2" + bl_options = {"REGISTER"} + + @classmethod + def poll(cls, context): + from .operators import active_job + + return is_image_empty(context.object) and active_job() is None + + def execute(self, context): + from .operators import IMAGE_BACKGROUND_JOB + + source_object = context.object + try: + source_path, delete_input = prepare_image_input( + source_object.data, + context.scene, + ) + except RuntimeError as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + + result = bpy.ops.snapform.run_job( + job_type=IMAGE_BACKGROUND_JOB, + snapform_input_path=str(source_path), + snapform_source_object_name=source_object.name, + snapform_delete_input=delete_input, + ) + if "CANCELLED" not in result: + return {"FINISHED"} + if delete_input: + source_path.unlink(missing_ok=True) + return {"CANCELLED"} + + class ConvertImageToDepthPlane(bpy.types.Operator): bl_idname = "snapform.convert_image_to_depth_plane" bl_label = "Convert to Depth Plane" @@ -739,6 +978,16 @@ class ConvertImageToDepthPlane(bpy.types.Operator): ), default="MODIFIER", ) + snapform_unshaded: bpy.props.BoolProperty( + name="Unshaded", + 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_model: bpy.props.EnumProperty( name="Model", description="Model used to estimate image depth", @@ -800,15 +1049,20 @@ class ConvertImageToDepthPlane(bpy.types.Operator): "snapform_process_res_factor", text=f"Resolution Scale ({process_res} px)", ) + layout.prop(self, "snapform_remove_background") layout.prop(self, "snapform_displacement_via") + layout.prop(self, "snapform_unshaded") layout.prop(self, "snapform_subdivision_level") layout.prop(self, "snapform_delete_original") def execute(self, context): - from .operators import IMAGE_DEPTH_JOB + from .operators import IMAGE_DEPTH_JOB, ORIGINAL_BACKGROUND_JOB source_object = context.object - if self.snapform_depth_source == "ORIGINAL": + if ( + self.snapform_depth_source == "ORIGINAL" + and not self.snapform_remove_background + ): try: build_image_depth_plane( context, @@ -816,6 +1070,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator): self.snapform_subdivision_level, self.snapform_delete_original, self.snapform_displacement_via, + self.snapform_unshaded, ) except (OSError, RuntimeError, ValueError) as error: self.report({"ERROR"}, str(error)) @@ -831,6 +1086,31 @@ class ConvertImageToDepthPlane(bpy.types.Operator): return {"CANCELLED"} return {"FINISHED"} + if self.snapform_depth_source == "ORIGINAL": + try: + source_path, delete_input = prepare_image_input( + source_object.data, + context.scene, + ) + except RuntimeError as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + result = bpy.ops.snapform.run_job( + job_type=ORIGINAL_BACKGROUND_JOB, + snapform_input_path=str(source_path), + snapform_source_object_name=source_object.name, + snapform_subdivision_level=self.snapform_subdivision_level, + snapform_delete_source=self.snapform_delete_original, + snapform_delete_input=delete_input, + snapform_displacement_via=self.snapform_displacement_via, + snapform_unshaded=self.snapform_unshaded, + ) + if "CANCELLED" not in result: + return {"FINISHED"} + if delete_input: + source_path.unlink(missing_ok=True) + return {"CANCELLED"} + image = source_object.data try: process_res = scaled_process_resolution( @@ -840,26 +1120,20 @@ class ConvertImageToDepthPlane(bpy.types.Operator): except ValueError as error: self.report({"ERROR"}, str(error)) return {"CANCELLED"} - source_path = image_path(image) - delete_input = False - if image_is_packed(image): - try: - source_path = save_packed_image(image, context.scene) - except RuntimeError as error: - self.report({"ERROR"}, str(error)) - return {"CANCELLED"} - delete_input = True - if not source_path.is_file(): - self.report( - {"ERROR"}, - "The image empty must reference an existing image file", + try: + source_path, delete_input = prepare_image_input( + image, + context.scene, ) + except RuntimeError as error: + self.report({"ERROR"}, str(error)) return {"CANCELLED"} result = bpy.ops.snapform.run_job( job_type=IMAGE_DEPTH_JOB, snapform_input_path=str(source_path), snapform_source_object_name=source_object.name, + snapform_depth_data_name=depth_data_name(source_object), snapform_model_key=self.snapform_model, snapform_use_ray_pose=self.snapform_use_ray_pose, snapform_process_res=process_res, @@ -867,6 +1141,8 @@ class ConvertImageToDepthPlane(bpy.types.Operator): snapform_delete_source=self.snapform_delete_original, snapform_delete_input=delete_input, snapform_displacement_via=self.snapform_displacement_via, + snapform_unshaded=self.snapform_unshaded, + snapform_remove_background=self.snapform_remove_background, ) if "CANCELLED" in result: if delete_input: @@ -889,3 +1165,10 @@ def draw_image_context_menu(self, context): text="Convert to Depth Plane", icon="MOD_DISPLACE", ) + row = self.layout.row() + row.operator_context = "INVOKE_DEFAULT" + row.operator( + "snapform.remove_image_background", + text="Remove Background", + icon="IMAGE_ALPHA", + ) diff --git a/src/snapform/models.py b/src/snapform/models.py index a4d6d72..2f438bf 100644 --- a/src/snapform/models.py +++ b/src/snapform/models.py @@ -9,6 +9,7 @@ class ModelSpec: license_name: str description: str supports_gs: bool + download_patterns: tuple[str, ...] = () MODELS = { @@ -53,6 +54,17 @@ MODELS = { supports_gs=True, ), } + +BEN2_MODEL = ModelSpec( + key="BEN2_BASE", + label="BEN2 Base", + repository="PramaLLC/BEN2", + license_name="MIT", + description="Fast high-quality background removal model", + supports_gs=False, + download_patterns=("config.json", "model.safetensors"), +) +DOWNLOADABLE_MODELS = {**MODELS, BEN2_MODEL.key: BEN2_MODEL} MODEL_NUMBERS = { model_key: number for number, model_key in enumerate(MODELS) } @@ -123,3 +135,10 @@ def get_model(model_key): return MODELS[model_key] except KeyError as error: raise ValueError(f"Unknown DA3 model: {model_key}") from error + + +def get_downloadable_model(model_key): + try: + return DOWNLOADABLE_MODELS[model_key] + except KeyError as error: + raise ValueError(f"Unknown model: {model_key}") from error diff --git a/src/snapform/operators.py b/src/snapform/operators.py index 1d2e870..3c62c38 100644 --- a/src/snapform/operators.py +++ b/src/snapform/operators.py @@ -5,11 +5,19 @@ from pathlib import Path import bpy from . import runtime -from .models import get_model, selected_model_key +from .models import ( + BEN2_MODEL, + get_downloadable_model, + get_model, + selected_model_key, +) RECONSTRUCTION_JOB = "RECONSTRUCT" IMAGE_DEPTH_JOB = "IMAGE_DEPTH" +BACKGROUND_REMOVAL_JOB = "REMOVE_BACKGROUND" +IMAGE_BACKGROUND_JOB = "IMAGE_BACKGROUND" +ORIGINAL_BACKGROUND_JOB = "ORIGINAL_BACKGROUND" IMAGE_DEPTH_DEFAULT_PROCESS_RES = 1024 IMAGE_DEPTH_PROCESS_RES_METHOD = "upper_bound_resize" @@ -33,10 +41,13 @@ class ActiveJob: result_directory: Path export_format: str = "" source_object_name: str = "" + depth_data_name: str = "" process_res: int = IMAGE_DEPTH_DEFAULT_PROCESS_RES subdivision_level: int = 5 delete_source: bool = True displacement_via: str = "MODIFIER" + unshaded: bool = False + remove_background: bool = False temporary_input_path: Path | None = None cancelled: bool = False @@ -83,6 +94,9 @@ class RunJob(bpy.types.Operator): ("DOWNLOAD_MODEL", "Download Model", ""), (RECONSTRUCTION_JOB, "Export DA3 Result", ""), (IMAGE_DEPTH_JOB, "Convert Image to Depth Plane", ""), + (BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""), + (IMAGE_BACKGROUND_JOB, "Remove Image Empty Background", ""), + (ORIGINAL_BACKGROUND_JOB, "Build Original Background Plane", ""), ), options={"HIDDEN", "SKIP_SAVE"}, ) @@ -95,6 +109,9 @@ class RunJob(bpy.types.Operator): snapform_source_object_name: bpy.props.StringProperty( options={"HIDDEN", "SKIP_SAVE"}, ) + snapform_depth_data_name: bpy.props.StringProperty( + options={"HIDDEN", "SKIP_SAVE"}, + ) snapform_model_key: bpy.props.StringProperty( default="BASE", options={"HIDDEN", "SKIP_SAVE"}, @@ -131,6 +148,14 @@ class RunJob(bpy.types.Operator): default="MODIFIER", options={"HIDDEN", "SKIP_SAVE"}, ) + snapform_unshaded: bpy.props.BoolProperty( + default=False, + options={"HIDDEN", "SKIP_SAVE"}, + ) + snapform_remove_background: bpy.props.BoolProperty( + default=False, + options={"HIDDEN", "SKIP_SAVE"}, + ) def execute(self, context): global _active_job @@ -173,13 +198,20 @@ class RunJob(bpy.types.Operator): else "" ), source_object_name=self.snapform_source_object_name, + depth_data_name=self.snapform_depth_data_name, process_res=self.snapform_process_res, subdivision_level=self.snapform_subdivision_level, delete_source=self.snapform_delete_source, displacement_via=self.snapform_displacement_via, + unshaded=self.snapform_unshaded, + remove_background=self.snapform_remove_background, temporary_input_path=( Path(self.snapform_input_path) - if self.job_type == IMAGE_DEPTH_JOB + if self.job_type in { + IMAGE_DEPTH_JOB, + IMAGE_BACKGROUND_JOB, + ORIGINAL_BACKGROUND_JOB, + } and self.snapform_delete_input else None ), @@ -196,6 +228,9 @@ class RunJob(bpy.types.Operator): def _prepare_command(self, _context, settings): paths = runtime.new_job_paths(self.job_type) + remove_background = bool( + getattr(self, "snapform_remove_background", False) + ) if self.job_type == "INSTALL_RUNTIME": self._require_online_access() return runtime.install_runtime_command(paths["status"]), paths @@ -205,20 +240,49 @@ class RunJob(bpy.types.Operator): if self.job_type == "DOWNLOAD_MODEL": self._require_online_access() - model = get_model(self.model_key) + model = get_downloadable_model(self.model_key) return ( runtime.download_model_command(model.key, paths["status"]), paths, ) - configured_input = ( - self.snapform_input_path - if self.job_type == IMAGE_DEPTH_JOB - else settings.input_path - ) + if self.job_type in { + IMAGE_DEPTH_JOB, + IMAGE_BACKGROUND_JOB, + ORIGINAL_BACKGROUND_JOB, + }: + configured_input = self.snapform_input_path + elif self.job_type == BACKGROUND_REMOVAL_JOB: + configured_input = settings.background_input_path + else: + configured_input = settings.input_path input_path = Path(bpy.path.abspath(configured_input)) if not input_path.is_file(): raise ValueError("Select a valid image or video file") + if self.job_type in { + BACKGROUND_REMOVAL_JOB, + IMAGE_BACKGROUND_JOB, + ORIGINAL_BACKGROUND_JOB, + }: + if input_path.suffix.lower() not in { + ".jpg", + ".jpeg", + ".png", + ".webp", + }: + raise ValueError("BEN2 requires a JPG, JPEG, PNG, or WEBP image") + if not runtime.model_ready(BEN2_MODEL.key): + raise RuntimeError(f"Download {BEN2_MODEL.label} first") + return ( + runtime.remove_background_command( + settings, + input_path, + paths["status"], + paths["result"], + BEN2_MODEL.key, + ), + paths, + ) if self.job_type not in {RECONSTRUCTION_JOB, IMAGE_DEPTH_JOB}: raise RuntimeError(f"Unsupported task type: {self.job_type}") export_format = ( @@ -236,6 +300,12 @@ class RunJob(bpy.types.Operator): model = get_model(model_key) if not runtime.model_ready(model_key): raise RuntimeError(f"Download {model.label} first") + if ( + self.job_type == IMAGE_DEPTH_JOB + and remove_background + and not runtime.model_ready(BEN2_MODEL.key) + ): + raise RuntimeError(f"Download {BEN2_MODEL.label} first") if export_format == "gs_ply" and not model.supports_gs: raise RuntimeError( "PLY export requires a DA3 Giant or Nested Giant model" @@ -255,6 +325,10 @@ class RunJob(bpy.types.Operator): else None ), use_ray_pose=use_ray_pose, + remove_background=( + self.job_type == IMAGE_DEPTH_JOB + and remove_background + ), ), paths, ) @@ -274,8 +348,15 @@ class RunJob(bpy.types.Operator): return f"Starting {settings.export_format} export" if self.job_type == IMAGE_DEPTH_JOB: return "Starting depth estimation" + if self.job_type == BACKGROUND_REMOVAL_JOB: + return "Starting BEN2 background removal" + if self.job_type == IMAGE_BACKGROUND_JOB: + return "Starting Image Empty background removal" + if self.job_type == ORIGINAL_BACKGROUND_JOB: + return "Starting background removal" if self.job_type == "DOWNLOAD_MODEL": - return f"Starting {get_model(self.model_key).label} download" + model = get_downloadable_model(self.model_key) + return f"Starting {model.label} download" return labels[self.job_type] def modal(self, context, event): @@ -330,15 +411,18 @@ class RunJob(bpy.types.Operator): self._redraw_ui(context, force=True) return {"CANCELLED"} - if job.kind == RECONSTRUCTION_JOB: + if job.kind in {RECONSTRUCTION_JOB, BACKGROUND_REMOVAL_JOB}: try: open_job_directory(job.result_directory) except Exception as open_error: self.report({"WARNING"}, str(open_error)) - self.report( - {"INFO"}, - f"Exported {job.export_format} to {job.result_directory}", - ) + if job.kind == RECONSTRUCTION_JOB: + message = ( + f"Exported {job.export_format} to {job.result_directory}" + ) + else: + message = f"Removed background to {job.result_directory}" + self.report({"INFO"}, message) if job.kind == IMAGE_DEPTH_JOB: try: settings.status = "Building depth plane" @@ -353,6 +437,34 @@ class RunJob(bpy.types.Operator): self._redraw_ui(context, force=True) return {"CANCELLED"} status["message"] = "Depth plane is ready" + if job.kind == ORIGINAL_BACKGROUND_JOB: + try: + settings.status = "Building depth plane" + settings.progress = 0.96 + self._redraw_ui(context, force=True) + self._build_original_depth_plane(job) + except (OSError, RuntimeError, ValueError) as import_error: + settings.status = "Depth import failed" + settings.progress = 0.0 + self.report({"ERROR"}, str(import_error)) + close_active_job() + self._redraw_ui(context, force=True) + return {"CANCELLED"} + status["message"] = "Depth plane is ready" + if job.kind == IMAGE_BACKGROUND_JOB: + try: + settings.status = "Updating Image Empty" + settings.progress = 0.96 + self._redraw_ui(context, force=True) + self._apply_removed_background(job) + except (OSError, RuntimeError, ValueError) as import_error: + settings.status = "Image update failed" + settings.progress = 0.0 + self.report({"ERROR"}, str(import_error)) + close_active_job() + self._redraw_ui(context, force=True) + return {"CANCELLED"} + status["message"] = "Image background is removed" settings.status = status.get("message", "Task complete") settings.progress = 1.0 @@ -370,9 +482,16 @@ class RunJob(bpy.types.Operator): result = bpy.ops.snapform.build_depth_plane( snapform_source_object_name=job.source_object_name, snapform_npz_path=str(npz_path), + snapform_depth_data_name=job.depth_data_name, + snapform_background_path=( + str(job.result_directory / "foreground.png") + if job.remove_background + else "" + ), snapform_subdivision_level=job.subdivision_level, snapform_delete_source=job.delete_source, snapform_displacement_via=job.displacement_via, + snapform_unshaded=job.unshaded, ) if "FINISHED" not in result: raise RuntimeError("Unable to build the depth plane") @@ -380,6 +499,37 @@ class RunJob(bpy.types.Operator): if "FINISHED" not in undo_result: raise RuntimeError("Unable to create the depth plane Undo step") + def _build_original_depth_plane(self, job): + result = bpy.ops.snapform.build_original_depth_plane( + snapform_source_object_name=job.source_object_name, + snapform_background_path=str( + job.result_directory / "foreground.png" + ), + snapform_subdivision_level=job.subdivision_level, + snapform_delete_source=job.delete_source, + snapform_displacement_via=job.displacement_via, + snapform_unshaded=job.unshaded, + ) + if "FINISHED" not in result: + raise RuntimeError("Unable to build the depth plane") + undo_result = bpy.ops.ed.undo_push(message="Convert to Depth Plane") + if "FINISHED" not in undo_result: + raise RuntimeError("Unable to create the depth plane Undo step") + + def _apply_removed_background(self, job, push_undo=True): + output_path = job.result_directory / "foreground.png" + result = bpy.ops.snapform.apply_removed_background( + snapform_source_object_name=job.source_object_name, + snapform_output_path=str(output_path), + ) + if "FINISHED" not in result: + raise RuntimeError("Unable to update the Image Empty") + if not push_undo: + return + undo_result = bpy.ops.ed.undo_push(message="Remove Background") + if "FINISHED" not in undo_result: + raise RuntimeError("Unable to create the background removal Undo step") + def _remove_timer(self, context): timer = getattr(self, "_timer", None) if timer is None: diff --git a/src/snapform/panel.py b/src/snapform/panel.py index 445018e..65c61ce 100644 --- a/src/snapform/panel.py +++ b/src/snapform/panel.py @@ -1,6 +1,6 @@ import bpy -from .models import MODELS +from .models import BEN2_MODEL, MODELS from .operators import active_job from .runtime import model_ready, runtime_ready @@ -29,6 +29,7 @@ class RuntimeModelsPanel(bpy.types.Panel): for model in MODELS.values(): self._draw_model(layout, model, running) + self._draw_model(layout, BEN2_MODEL, running) def _draw_model(self, layout, model, running): row = layout.row(align=True) @@ -106,3 +107,15 @@ class DebugPanel(bpy.types.Panel): icon="EXPORT", ) operator.job_type = "RECONSTRUCT" + + background_box = layout.box() + background_box.label(text="BEN2 Background Removal") + background_box.prop(settings, "background_input_path") + background_row = background_box.row() + background_row.enabled = active_job() is None + operator = background_row.operator( + "snapform.run_job", + text="Remove Background", + icon="IMAGE_ALPHA", + ) + operator.job_type = "REMOVE_BACKGROUND" diff --git a/src/snapform/properties.py b/src/snapform/properties.py index 55f06ce..2b66aad 100644 --- a/src/snapform/properties.py +++ b/src/snapform/properties.py @@ -9,6 +9,11 @@ class SnapFormSettings(bpy.types.PropertyGroup): description="A single image or video file used for reconstruction", subtype="FILE_PATH", ) + background_input_path: bpy.props.StringProperty( + name="Image", + description="A single image used to test BEN2 background removal", + subtype="FILE_PATH", + ) video_frames: bpy.props.IntProperty( name="Video Frames", description="Maximum number of frames sampled evenly from a video", diff --git a/src/snapform/runtime.py b/src/snapform/runtime.py index da3219c..e608ff9 100644 --- a/src/snapform/runtime.py +++ b/src/snapform/runtime.py @@ -8,11 +8,12 @@ from pathlib import Path import bpy -from .models import get_model, selected_model_key +from .models import BEN2_MODEL, get_downloadable_model, selected_model_key DA3_COMMIT = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4" -RUNTIME_SCHEMA = 2 +BEN2_COMMIT = "2c99a5da477b5523585bfa5c893888a6e818a8f6" +RUNTIME_SCHEMA = 3 DEFAULT_STORAGE_ROOT = Path.home() / ".snapform" @@ -127,14 +128,16 @@ def install_runtime_command(status_path): backend.lower(), "--da3-commit", DA3_COMMIT, + "--ben2-commit", + BEN2_COMMIT, "--status", str(status_path), ] def download_model_command(model_key, status_path): - model = get_model(model_key) - return [ + model = get_downloadable_model(model_key) + command = [ str(runtime_python()), str(Path(__file__).with_name("worker.py")), "download", @@ -145,6 +148,33 @@ def download_model_command(model_key, status_path): "--status", str(status_path), ] + for pattern in model.download_patterns: + command.extend(["--allow-pattern", pattern]) + return command + + +def remove_background_command( + settings, + input_path, + status_path, + result_directory, + model_key, +): + return [ + str(runtime_python()), + str(Path(__file__).with_name("worker.py")), + "remove-background", + "--input", + str(input_path), + "--model-dir", + str(model_directory(model_key)), + "--output-dir", + str(result_directory), + "--status", + str(status_path), + "--device", + settings.device.lower(), + ] def reconstruct_command( @@ -157,6 +187,7 @@ def reconstruct_command( process_res=None, process_res_method=None, use_ray_pose=None, + remove_background=False, ): snapform_selected_key = model_key or selected_model_key(settings, export_format) snapform_selected_process_res = ( @@ -170,7 +201,7 @@ def reconstruct_command( snapform_selected_ray_pose = ( settings.use_ray_pose if use_ray_pose is None else use_ray_pose ) - return [ + command = [ str(runtime_python()), str(Path(__file__).with_name("worker.py")), "reconstruct", @@ -209,6 +240,14 @@ def reconstruct_command( "--ref-view-strategy", settings.ref_view_strategy, ] + if remove_background: + command.extend( + [ + "--remove-background-model-dir", + str(model_directory(BEN2_MODEL.key)), + ] + ) + return command def creation_flags(): diff --git a/src/snapform/worker.py b/src/snapform/worker.py index 9a41b39..49f202d 100644 --- a/src/snapform/worker.py +++ b/src/snapform/worker.py @@ -188,6 +188,7 @@ def download_model(args): repo_id=args.repository, local_dir=model_dir, tqdm_class=progress_class, + allow_patterns=args.allow_pattern or None, ) write_status(args.status, 1.0, "Model is ready") @@ -202,6 +203,76 @@ def select_device(requested, torch): return torch.device("cuda" if torch.cuda.is_available() else "cpu") +def infer_foreground(input_path, model_dir, device, torch): + from ben2 import BEN_Base + from PIL import Image + + input_path = Path(input_path) + if input_path.suffix.lower() not in IMAGE_SUFFIXES: + raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image") + load_started = time.perf_counter() + model = BEN_Base.from_pretrained(model_dir) + model = model.to(device=device) + model.eval() + load_ms = (time.perf_counter() - load_started) * 1000.0 + + with Image.open(input_path) as opened_image: + image = opened_image.convert("RGB") + if device.type == "cuda": + torch.cuda.synchronize() + inference_started = time.perf_counter() + with torch.inference_mode(): + foreground = model.inference( + image, + refine_foreground=False, + ) + if device.type == "cuda": + torch.cuda.synchronize() + inference_ms = (time.perf_counter() - inference_started) * 1000.0 + del model + if device.type == "cuda": + torch.cuda.empty_cache() + return foreground, load_ms, inference_ms + + +def remove_background(args): + total_started = time.perf_counter() + write_status(args.status, 0.02, "Starting BEN2 Runtime") + + import torch + + output_directory = Path(args.output_dir) + output_directory.mkdir(parents=True, exist_ok=False) + device = select_device(args.device, torch) + + write_status(args.status, 0.12, f"Loading BEN2 on {device.type}") + write_status(args.status, 0.35, "Removing background") + foreground, load_ms, inference_ms = infer_foreground( + args.input, + args.model_dir, + device, + torch, + ) + + write_status(args.status, 0.92, "Writing transparent PNG") + output_path = output_directory / "foreground.png" + foreground.convert("RGBA").save( + output_path, + format="PNG", + compress_level=1, + ) + if not output_path.is_file(): + raise RuntimeError("BEN2 did not write a foreground image") + + total_ms = (time.perf_counter() - total_started) * 1000.0 + message = ( + f"Background removed in {inference_ms:.0f} ms " + f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)" + ) + print(message) + write_status(args.status, 1.0, message) + + def exported_paths(output_directory, export_format): output_directory = Path(output_directory) if export_format == "glb": @@ -246,6 +317,30 @@ def reconstruct(args): try: device = select_device(args.device, torch) + remove_background_model_dir = getattr( + args, + "remove_background_model_dir", + "", + ) + if remove_background_model_dir: + write_status(args.status, 0.12, "Removing image background") + foreground, load_ms, inference_ms = infer_foreground( + images[0], + remove_background_model_dir, + device, + torch, + ) + foreground.convert("RGBA").save( + output_directory / "foreground.png", + format="PNG", + compress_level=1, + ) + if not (output_directory / "foreground.png").is_file(): + raise RuntimeError("BEN2 did not write a foreground image") + print( + f"BEN2 inference {inference_ms:.0f} ms; " + f"model load {load_ms:.0f} ms" + ) write_status(args.status, 0.15, f"Loading model on {device.type}") model = DepthAnything3.from_pretrained(args.model_dir) model = model.to(device=device) @@ -311,6 +406,18 @@ def parser(): download.add_argument("--repository", required=True) download.add_argument("--model-dir", required=True) download.add_argument("--status", required=True) + download.add_argument("--allow-pattern", action="append", default=[]) + + background = commands.add_parser("remove-background") + background.add_argument("--input", required=True) + background.add_argument("--model-dir", required=True) + background.add_argument("--output-dir", required=True) + background.add_argument("--status", required=True) + background.add_argument( + "--device", + choices=("auto", "cuda", "cpu"), + required=True, + ) reconstruction = commands.add_parser("reconstruct") reconstruction.add_argument("--input", required=True) @@ -322,6 +429,7 @@ def parser(): required=True, ) reconstruction.add_argument("--status", required=True) + reconstruction.add_argument("--remove-background-model-dir", default="") reconstruction.add_argument("--process-res", type=int, required=True) reconstruction.add_argument( "--process-res-method", @@ -370,8 +478,10 @@ def main(): try: if args.command == "download": download_model(args) - else: + elif args.command == "reconstruct": reconstruct(args) + else: + remove_background(args) except Exception as error: write_status(args.status, 1.0, "Task failed", str(error)) raise diff --git a/tests/test_addon.py b/tests/test_addon.py index 5985fe5..2135985 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -118,12 +118,26 @@ class AddonTest(unittest.TestCase): ) self.assertFalse(hasattr(self.addon, "ImportDepthPlane")) + def test_registers_image_background_operators(self): + self.assertIn(self.addon.RemoveImageBackground, self.addon.CLASSES) + self.assertIn(self.addon.ApplyRemovedBackground, self.addon.CLASSES) + self.assertIn(self.addon.BuildOriginalDepthPlane, self.addon.CLASSES) + self.assertEqual( + self.addon.ApplyRemovedBackground.bl_options, + {"INTERNAL"}, + ) + def test_completed_depth_job_uses_undoable_build_operator(self): source = inspect.getsource( self.addon.RunJob._build_depth_plane ) self.assertIn("bpy.ops.snapform.build_depth_plane(", source) + self.assertIn("snapform_unshaded=job.unshaded", source) + self.assertIn( + "snapform_depth_data_name=job.depth_data_name", + source, + ) self.assertIn('if "FINISHED" not in result:', source) self.assertIn( 'bpy.ops.ed.undo_push(message="Convert to Depth Plane")', @@ -131,7 +145,7 @@ class AddonTest(unittest.TestCase): ) self.assertIn('if "FINISHED" not in undo_result:', source) - def test_image_context_menu_only_exposes_conversion(self): + def test_image_context_menu_exposes_depth_and_background_actions(self): draw_source = inspect.getsource( self.addon.draw_image_context_menu ) @@ -140,10 +154,185 @@ class AddonTest(unittest.TestCase): '"snapform.convert_image_to_depth_plane"', draw_source, ) + self.assertIn( + '"snapform.remove_image_background"', + draw_source, + ) + self.assertIn('text="Remove Background"', draw_source) + self.assertEqual(draw_source.count("row = self.layout.row()"), 2) self.assertIn('operator_context = "INVOKE_DEFAULT"', draw_source) self.assertNotIn("import_depth_plane", draw_source) self.assertNotIn("Rebuild", draw_source) + def test_background_result_creates_new_packed_image(self): + original_image = SimpleNamespace( + name="img.png", + use_fake_user=False, + ) + transform = object() + source_object = SimpleNamespace( + data=original_image, + matrix_world=transform, + empty_display_size=2.5, + empty_image_offset=(-0.25, -0.75), + ) + pack_calls = [] + result_image = SimpleNamespace( + name="foreground.png", + pack=lambda: pack_calls.append(True), + ) + load_calls = [] + remove_calls = [] + self.fake_bpy.data = SimpleNamespace( + images=SimpleNamespace( + load=lambda *args, **options: load_calls.append( + (args, options) + ) + or result_image, + remove=remove_calls.append, + ), + ) + + with TemporaryDirectory() as directory: + output_path = Path(directory) / "foreground.png" + output_path.touch() + created = self.addon.image_depth.create_removed_background_image( + source_object, + output_path, + ) + + self.assertIs(created, result_image) + self.assertIs(source_object.data, result_image) + self.assertIsNot(source_object.data, original_image) + self.assertTrue(original_image.use_fake_user) + self.assertIs(source_object.matrix_world, transform) + self.assertEqual(source_object.empty_display_size, 2.5) + self.assertEqual( + source_object.empty_image_offset, + (-0.25, -0.75), + ) + self.assertEqual(result_image.name, "img.rb") + self.assertEqual(pack_calls, [True]) + self.assertEqual(remove_calls, []) + self.assertEqual( + load_calls[0][1], + {"check_existing": False}, + ) + + def test_background_result_name_handles_numbered_format_suffix(self): + source_object = SimpleNamespace( + data=SimpleNamespace(name="img.png.001"), + ) + + self.assertEqual( + self.addon.image_depth.removed_background_data_name( + source_object + ), + "img.001.rb", + ) + + def test_depth_background_image_does_not_replace_source_empty(self): + original_image = SimpleNamespace( + name="img.png", + use_fake_user=False, + ) + source_object = SimpleNamespace(data=original_image) + pack_calls = [] + result_image = SimpleNamespace( + name="foreground.png", + pack=lambda: pack_calls.append(True), + ) + self.fake_bpy.data = SimpleNamespace( + images=SimpleNamespace( + load=lambda *_args, **_options: result_image, + remove=lambda _image: None, + ), + ) + + with TemporaryDirectory() as directory: + output_path = Path(directory) / "foreground.png" + output_path.touch() + created = self.addon.image_depth.create_removed_background_image( + source_object, + output_path, + assign_to_source=False, + ) + + self.assertIs(created, result_image) + self.assertIs(source_object.data, original_image) + self.assertFalse(original_image.use_fake_user) + self.assertEqual(result_image.name, "img.rb") + self.assertEqual(pack_calls, [True]) + + def test_image_empty_background_action_starts_production_job(self): + source_object = SimpleNamespace( + name="Reference", + data=SimpleNamespace(name="img.png"), + ) + context = SimpleNamespace(object=source_object, scene=object()) + calls = [] + self.fake_bpy.ops = SimpleNamespace( + snapform=SimpleNamespace( + run_job=lambda **options: calls.append(options) + or {"RUNNING_MODAL"}, + ), + ) + operator = self.addon.RemoveImageBackground() + + with TemporaryDirectory() as directory: + input_path = Path(directory) / "input.png" + input_path.touch() + with patch.object( + self.addon.image_depth, + "prepare_image_input", + return_value=(input_path, True), + ): + result = operator.execute(context) + + self.assertEqual(result, {"FINISHED"}) + self.assertEqual(calls[0]["job_type"], "IMAGE_BACKGROUND") + self.assertEqual(calls[0]["snapform_input_path"], str(input_path)) + self.assertEqual( + calls[0]["snapform_source_object_name"], + "Reference", + ) + self.assertTrue(calls[0]["snapform_delete_input"]) + + 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_displacement_via = "MODIFIER" + operator.snapform_unshaded = False + operator.snapform_subdivision_level = 5 + operator.snapform_delete_original = False + source_object = SimpleNamespace( + name="Reference", + data=SimpleNamespace(name="img.png"), + ) + context = SimpleNamespace(object=source_object, scene=object()) + calls = [] + self.fake_bpy.ops = SimpleNamespace( + snapform=SimpleNamespace( + run_job=lambda **options: calls.append(options) + or {"RUNNING_MODAL"}, + ), + ) + + with TemporaryDirectory() as directory: + input_path = Path(directory) / "input.png" + input_path.touch() + with patch.object( + self.addon.image_depth, + "prepare_image_input", + return_value=(input_path, False), + ): + result = operator.execute(context) + + self.assertEqual(result, {"FINISHED"}) + self.assertEqual(calls[0]["job_type"], "ORIGINAL_BACKGROUND") + self.assertFalse(calls[0]["snapform_delete_source"]) + def test_conversion_dialog_exposes_requested_settings(self): operator = self.addon.ConvertImageToDepthPlane @@ -159,6 +348,12 @@ class AddonTest(unittest.TestCase): self.assertTrue( operator.__annotations__["snapform_delete_original"]["default"] ) + self.assertFalse( + operator.__annotations__["snapform_unshaded"]["default"] + ) + self.assertFalse( + operator.__annotations__["snapform_remove_background"]["default"] + ) self.assertEqual( operator.__annotations__["snapform_subdivision_level"]["default"], 5, @@ -168,6 +363,8 @@ class AddonTest(unittest.TestCase): { "snapform_depth_source", "snapform_displacement_via", + "snapform_unshaded", + "snapform_remove_background", "snapform_model", "snapform_use_ray_pose", "snapform_process_res_factor", @@ -192,6 +389,10 @@ class AddonTest(unittest.TestCase): [item[0] for item in displacement_via["items"]], ["MODIFIER", "MATERIAL"], ) + self.assertEqual( + operator.__annotations__["snapform_unshaded"]["name"], + "Unshaded", + ) draw_source = inspect.getsource(operator.draw) self.assertIn( 'if self.snapform_depth_source == "ESTIMATE":', @@ -218,8 +419,16 @@ class AddonTest(unittest.TestCase): 'estimate_box.prop(self, "snapform_use_ray_pose")', draw_source, ) + self.assertIn( + 'layout.prop(self, "snapform_remove_background")', + draw_source, + ) self.assertNotIn("layout.separator", draw_source) self.assertNotIn("layout.label", draw_source) + self.assertIn( + 'layout.prop(self, "snapform_unshaded")', + draw_source, + ) model_items = operator.__annotations__["snapform_model"]["items"] model_items = model_items(None, None) self.assertEqual( @@ -266,6 +475,12 @@ class AddonTest(unittest.TestCase): self.assertNotIn("DA3 image depth", starting_source) self.assertNotIn("DA3 depth plane", finish_source) + def test_depth_background_is_loaded_without_changing_source_empty(self): + source = inspect.getsource(self.addon.RunJob._build_depth_plane) + + self.assertIn("snapform_background_path=", source) + self.assertNotIn("_apply_removed_background", source) + def test_model_options_mark_only_models_that_are_not_downloaded(self): items_callback = self.addon.SnapFormSettings.__annotations__["model"][ "items" @@ -288,9 +503,12 @@ class AddonTest(unittest.TestCase): operator.snapform_use_ray_pose = True operator.snapform_process_res_factor = 0.5 operator.snapform_displacement_via = "MATERIAL" + operator.snapform_unshaded = True + operator.snapform_remove_background = True operator.snapform_subdivision_level = 5 operator.snapform_delete_original = True image = SimpleNamespace( + name="img.png", filepath="", packed_file=object(), size=(1920, 1080), @@ -326,6 +544,9 @@ class AddonTest(unittest.TestCase): calls[0]["snapform_displacement_via"], "MATERIAL", ) + self.assertTrue(calls[0]["snapform_unshaded"]) + self.assertTrue(calls[0]["snapform_remove_background"]) + self.assertEqual(calls[0]["snapform_depth_data_name"], "img.depth") self.assertTrue(calls[0]["snapform_delete_input"]) self.assertEqual( self.addon.image_depth._last_depth_plane_model, @@ -454,6 +675,81 @@ class AddonTest(unittest.TestCase): snapform_ray_pose, ) + def test_image_depth_background_option_requires_ben2(self): + settings = SimpleNamespace( + input_path="debug-input.png", + export_format="mini_npz", + model="BASE", + gs_model="GIANT_1_1", + ) + operator = self.addon.RunJob() + operator.job_type = self.addon.operators.IMAGE_DEPTH_JOB + operator.snapform_model_key = "BASE" + operator.snapform_use_ray_pose = True + operator.snapform_process_res = 1024 + operator.snapform_remove_background = True + + with TemporaryDirectory() as directory: + input_path = Path(directory) / "input.png" + input_path.touch() + operator.snapform_input_path = str(input_path) + with patch.object( + self.addon.runtime, + "new_job_paths", + return_value={ + "log": Path(directory) / "job.log", + "status": Path(directory) / "job.json", + "result": Path(directory) / "result", + }, + ), patch.object( + self.addon.runtime, + "runtime_ready", + return_value=True, + ), patch.object( + self.addon.runtime, + "model_ready", + side_effect=lambda key: key == "BASE", + ): + with self.assertRaisesRegex(RuntimeError, "BEN2 Base"): + operator._prepare_command(None, settings) + + def test_image_depth_background_option_reaches_worker_command(self): + runtime = importlib.import_module("snapform.runtime") + settings = SimpleNamespace( + input_path="input.png", + model="BASE", + gs_model="GIANT_1_1", + process_res=504, + process_res_method="upper_bound_resize", + conf_thresh_percentile=40.0, + num_max_points=1_000_000, + show_cameras=True, + gs_views_interval=1, + video_frames=16, + device="AUTO", + use_ray_pose=False, + ref_view_strategy="saddle_balanced", + ) + with patch.object( + runtime, + "runtime_python", + return_value=Path("python"), + ), patch.object( + runtime, + "model_directory", + side_effect=lambda key: Path(key.lower()), + ): + command = runtime.reconstruct_command( + settings, + "mini_npz", + Path("status.json"), + Path("result"), + remove_background=True, + ) + + option_index = command.index("--remove-background-model-dir") + self.assertEqual(command[option_index + 1], str(Path("ben2_base"))) + def test_preferences_use_package_identifier(self): self.assertEqual( self.addon.SnapFormPreferences.bl_idname, @@ -625,6 +921,13 @@ class AddonTest(unittest.TestCase): self.assertNotIn('text="Missing"', draw_source) self.assertNotIn("License:", draw_source) self.assertIn('operator.job_type = "RECONSTRUCT"', draw_source) + self.assertIn('text="BEN2 Background Removal"', draw_source) + self.assertIn( + 'background_box.prop(settings, "background_input_path")', + draw_source, + ) + self.assertIn('text="Remove Background"', draw_source) + self.assertIn('operator.job_type = "REMOVE_BACKGROUND"', draw_source) def test_runtime_models_panel_assigns_each_download_target(self): draw_source = inspect.getsource( @@ -635,6 +938,10 @@ class AddonTest(unittest.TestCase): ) self.assertIn("for model in MODELS.values()", draw_source) + self.assertIn( + "self._draw_model(layout, BEN2_MODEL, running)", + draw_source, + ) self.assertIn("operator.model_key = model.key", model_source) self.assertIn( "download_row.enabled = runtime_ready() and not running", @@ -648,6 +955,191 @@ class AddonTest(unittest.TestCase): "Setup", ) + def test_ben2_debug_job_uses_its_image_and_model(self): + settings = SimpleNamespace( + background_input_path="input.png", + device="CUDA", + ) + operator = self.addon.RunJob() + operator.job_type = self.addon.operators.BACKGROUND_REMOVAL_JOB + + with TemporaryDirectory() as directory: + root = Path(directory) + input_path = root / "input.png" + input_path.touch() + settings.background_input_path = str(input_path) + paths = { + "log": root / "job.log", + "status": root / "job.json", + "result": root / "result", + } + with patch.object( + self.addon.runtime, + "new_job_paths", + return_value=paths, + ), patch.object( + self.addon.runtime, + "runtime_ready", + return_value=True, + ), patch.object( + self.addon.runtime, + "model_ready", + return_value=True, + ) as model_ready, patch.object( + self.addon.runtime, + "remove_background_command", + return_value=["worker"], + ) as remove_background: + command, prepared_paths = operator._prepare_command( + None, + settings, + ) + + self.assertEqual(command, ["worker"]) + self.assertEqual(prepared_paths, paths) + model_ready.assert_called_once_with("BEN2_BASE") + self.assertEqual(remove_background.call_args.args[1], input_path) + self.assertEqual(remove_background.call_args.args[4], "BEN2_BASE") + + def test_image_background_job_uses_operator_input(self): + settings = SimpleNamespace(device="AUTO") + operator = self.addon.RunJob() + operator.job_type = self.addon.operators.IMAGE_BACKGROUND_JOB + + with TemporaryDirectory() as directory: + root = Path(directory) + input_path = root / "input.png" + input_path.touch() + operator.snapform_input_path = str(input_path) + paths = { + "log": root / "job.log", + "status": root / "job.json", + "result": root / "result", + } + with patch.object( + self.addon.runtime, + "new_job_paths", + return_value=paths, + ), patch.object( + self.addon.runtime, + "runtime_ready", + return_value=True, + ), patch.object( + self.addon.runtime, + "model_ready", + return_value=True, + ), patch.object( + self.addon.runtime, + "remove_background_command", + return_value=["worker"], + ) as remove_background: + command, prepared_paths = operator._prepare_command( + None, + settings, + ) + + self.assertEqual(command, ["worker"]) + self.assertEqual(prepared_paths, paths) + self.assertEqual(remove_background.call_args.args[1], input_path) + + def test_completed_image_background_job_applies_result_with_undo(self): + calls = [] + self.fake_bpy.ops = SimpleNamespace( + snapform=SimpleNamespace( + apply_removed_background=lambda **options: calls.append( + ("apply", options) + ) + or {"FINISHED"}, + ), + ed=SimpleNamespace( + undo_push=lambda **options: calls.append(("undo", options)) + or {"FINISHED"}, + ), + ) + operator = self.addon.RunJob() + job = SimpleNamespace( + result_directory=Path("result"), + source_object_name="Reference", + ) + + operator._apply_removed_background(job) + + self.assertEqual( + calls, + [ + ( + "apply", + { + "snapform_source_object_name": "Reference", + "snapform_output_path": str( + Path("result") / "foreground.png" + ), + }, + ), + ("undo", {"message": "Remove Background"}), + ], + ) + + def test_completed_original_background_job_builds_with_one_undo(self): + calls = [] + self.fake_bpy.ops = SimpleNamespace( + snapform=SimpleNamespace( + build_original_depth_plane=lambda **options: calls.append( + ("build", options) + ) + or {"FINISHED"}, + ), + ed=SimpleNamespace( + undo_push=lambda **options: calls.append(("undo", options)) + or {"FINISHED"}, + ), + ) + operator = self.addon.RunJob() + job = SimpleNamespace( + result_directory=Path("result"), + source_object_name="Reference", + subdivision_level=5, + delete_source=False, + displacement_via="MODIFIER", + unshaded=False, + ) + + operator._build_original_depth_plane(job) + + self.assertEqual(calls[0][0], "build") + self.assertEqual( + calls[0][1]["snapform_background_path"], + str(Path("result") / "foreground.png"), + ) + self.assertFalse(calls[0][1]["snapform_delete_source"]) + self.assertEqual( + calls[1], + ("undo", {"message": "Convert to Depth Plane"}), + ) + + def test_ben2_download_command_filters_large_alternative_weights(self): + runtime = importlib.import_module("snapform.runtime") + with patch.object( + runtime, + "runtime_python", + return_value=Path("python"), + ), patch.object( + runtime, + "model_directory", + return_value=Path("model"), + ): + command = runtime.download_model_command( + "BEN2_BASE", + Path("status.json"), + ) + + patterns = [ + command[index + 1] + for index, value in enumerate(command) + if value == "--allow-pattern" + ] + self.assertEqual(patterns, ["config.json", "model.safetensors"]) + def test_preferences_install_button_is_disabled_while_busy(self): draw_source = inspect.getsource(self.addon.SnapFormPreferences.draw) diff --git a/tests/test_bootstrap_runtime.py b/tests/test_bootstrap_runtime.py index 2c94504..f390541 100644 --- a/tests/test_bootstrap_runtime.py +++ b/tests/test_bootstrap_runtime.py @@ -31,6 +31,7 @@ class BootstrapRuntimeTest(unittest.TestCase): runtime_dir=str(runtime_dir), backend="cu128", da3_commit="test-commit", + ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object(bootstrap, "run") as mocked_run: @@ -58,6 +59,22 @@ class BootstrapRuntimeTest(unittest.TestCase): ) self.assertEqual(gsplat_call.args[1], {"BUILD_NO_CUDA": "1"}) + ben2_command = next( + command + for command in commands + if any("ben2 @" in str(value) for value in command) + ) + self.assertTrue( + any("test-ben2-commit" in str(value) for value in ben2_command) + ) + check_command = commands[-1] + self.assertTrue( + any( + "from ben2 import BEN_Base" in str(value) + for value in check_command + ) + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_image_depth.py b/tests/test_image_depth.py index 0ee0b32..2b9ba92 100644 --- a/tests/test_image_depth.py +++ b/tests/test_image_depth.py @@ -298,9 +298,52 @@ class ImageDepthTest(unittest.TestCase): original_source = inspect.getsource( self.image_depth.build_image_depth_plane ) - self.assertIn("depth_name = depth_data_name(source_object)", estimate_source) + self.assertIn( + "depth_name = depth_name or depth_data_name(source_object)", + estimate_source, + ) self.assertIn("depth_data_name(source_object)", original_source) + def test_removed_background_name_uses_source_image_base_name(self): + source_object = SimpleNamespace( + data=SimpleNamespace(name="img.png"), + ) + + self.assertEqual( + self.image_depth.removed_background_data_name(source_object), + "img.rb", + ) + source_object.data.name = "img.png.001" + self.assertEqual( + self.image_depth.removed_background_data_name(source_object), + "img.001.rb", + ) + + def test_depth_name_override_keeps_original_base_name(self): + source = inspect.getsource(self.image_depth.build_depth_plane) + + self.assertIn( + "depth_name = depth_name or depth_data_name(source_object)", + source, + ) + self.assertIn( + "color_image if color_image is not None else source_object.data", + source, + ) + + def test_original_depth_background_only_changes_material_color(self): + source = inspect.getsource(self.image_depth.build_image_depth_plane) + + self.assertIn("texture.image = source_object.data", source) + self.assertIn( + "source_object.data if displacement_via == \"MATERIAL\" else None", + source, + ) + self.assertIn( + "color_image if color_image is not None else source_object.data", + source, + ) + def test_material_uses_source_object_name(self): da3_source = inspect.getsource(self.image_depth.build_depth_plane) image_source = inspect.getsource( @@ -319,7 +362,10 @@ class ImageDepthTest(unittest.TestCase): self.image_depth.build_depth_plane ) - self.assertIn("depth_name = depth_data_name(source_object)", source) + self.assertIn( + "depth_name = depth_name or depth_data_name(source_object)", + source, + ) self.assertIn("create_depth_image(\n depth_name,", source) self.assertIn("bpy.data.textures.new(\n depth_name,", source) self.assertIn('texture.extension = "EXTEND"', source) @@ -385,6 +431,42 @@ class ImageDepthTest(unittest.TestCase): ) self.assertIn("set_displacement_only(material)", source) + def test_unshaded_material_uses_alpha_mixed_emission(self): + source = inspect.getsource(self.image_depth.create_material) + + unshaded_start = source.index("if unshaded:") + shaded_start = source.index("else:", unshaded_start) + displacement_start = source.index( + "if displacement_image is not None:", + shaded_start, + ) + unshaded_source = source[unshaded_start:shaded_start] + shaded_source = source[shaded_start:displacement_start] + + self.assertIn('nodes.new("ShaderNodeEmission")', unshaded_source) + self.assertIn( + 'nodes.new("ShaderNodeBsdfTransparent")', + unshaded_source, + ) + self.assertIn('nodes.new("ShaderNodeMixShader")', unshaded_source) + self.assertIn( + 'color_texture.outputs["Color"]', + unshaded_source, + ) + self.assertIn('emission.inputs["Color"]', unshaded_source) + self.assertIn('color_texture.outputs["Alpha"]', unshaded_source) + self.assertIn('mix_shader.inputs["Fac"]', unshaded_source) + self.assertIn('transparent.outputs["BSDF"]', unshaded_source) + self.assertIn('mix_shader.inputs[1]', unshaded_source) + self.assertIn('emission.outputs["Emission"]', unshaded_source) + self.assertIn('mix_shader.inputs[2]', unshaded_source) + self.assertIn('mix_shader.outputs["Shader"]', unshaded_source) + self.assertNotIn("ShaderNodeBsdfPrincipled", unshaded_source) + self.assertNotIn("ShaderNodeBump", unshaded_source) + self.assertNotIn("ShaderNodeFloatCurve", unshaded_source) + self.assertIn("ShaderNodeBsdfPrincipled", shaded_source) + self.assertGreater(displacement_start, unshaded_start) + def test_displacement_values_follow_depth_source(self): estimate_source = inspect.getsource( self.image_depth.build_depth_plane @@ -434,6 +516,7 @@ class ImageDepthTest(unittest.TestCase): 'if displacement_via == "MATERIAL" else None', source, ) + self.assertIn("unshaded=unshaded", source) def test_surface_displacement_is_set_to_displacement_only(self): class Material: diff --git a/tests/test_models.py b/tests/test_models.py index fc02145..5b9f794 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -43,3 +43,14 @@ class ModelTest(unittest.TestCase): def test_unknown_model_is_rejected(self): with self.assertRaisesRegex(ValueError, "Unknown DA3 model"): self.models.get_model("UNKNOWN") + + def test_ben2_download_only_includes_required_hugging_face_files(self): + model = self.models.get_downloadable_model("BEN2_BASE") + + self.assertEqual(model.repository, "PramaLLC/BEN2") + self.assertEqual(model.license_name, "MIT") + self.assertEqual( + model.download_patterns, + ("config.json", "model.safetensors"), + ) + self.assertNotIn("BEN2_BASE", self.models.MODELS) diff --git a/tests/test_pack.py b/tests/test_pack.py index bac218c..7f3b079 100644 --- a/tests/test_pack.py +++ b/tests/test_pack.py @@ -29,7 +29,7 @@ class PackTest(unittest.TestCase): finally: pack.DIST_DIR = original_dist - self.assertEqual(output_path.name, "SnapForm.v0.9.27.zip") + self.assertEqual(output_path.name, "SnapForm.v0.9.28.zip") with ZipFile(output_path) as archive: names = set(archive.namelist()) diff --git a/tests/test_worker.py b/tests/test_worker.py index 17a1bab..1aaa7e9 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -47,6 +47,52 @@ class WorkerTest(unittest.TestCase): source.index("import torch"), ) + def test_optional_background_removal_runs_before_da3_model_load(self): + source = inspect.getsource(self.worker.reconstruct) + + self.assertLess( + source.index("infer_foreground("), + source.index("DepthAnything3.from_pretrained"), + ) + args = self.worker.parser().parse_args( + [ + "reconstruct", + "--input", + "input.png", + "--model-dir", + "da3", + "--output-dir", + "result", + "--export-format", + "mini_npz", + "--status", + "status.json", + "--remove-background-model-dir", + "ben2", + "--process-res", + "1024", + "--process-res-method", + "upper_bound_resize", + "--conf-thresh-percentile", + "40", + "--num-max-points", + "1000", + "--show-cameras", + "0", + "--gs-views-interval", + "1", + "--video-frames", + "16", + "--device", + "auto", + "--use-ray-pose", + "1", + "--ref-view-strategy", + "first", + ] + ) + self.assertEqual(args.remove_background_model_dir, "ben2") + def test_prepares_a_single_image_file(self): with tempfile.TemporaryDirectory() as directory: image_path = Path(directory) / "input.png" @@ -69,6 +115,117 @@ class WorkerTest(unittest.TestCase): with self.assertRaisesRegex(RuntimeError, "Unsupported input type"): self.worker.prepare_input(input_path, 16, directory) + def test_download_passes_optional_file_patterns(self): + calls = [] + fake_hugging_face = ModuleType("huggingface_hub") + fake_hugging_face.snapshot_download = lambda **options: calls.append( + options + ) + args = SimpleNamespace( + repository="PramaLLC/BEN2", + model_dir="model", + status="status.json", + allow_pattern=["config.json", "model.safetensors"], + ) + + with tempfile.TemporaryDirectory() as directory, patch.dict( + sys.modules, + {"huggingface_hub": fake_hugging_face}, + ), patch.object( + self.worker, + "create_download_progress_class", + return_value=object, + ): + args.model_dir = str(Path(directory) / "model") + args.status = str(Path(directory) / "status.json") + self.worker.download_model(args) + + self.assertEqual( + calls[0]["allow_patterns"], + ["config.json", "model.safetensors"], + ) + + def test_ben2_background_removal_writes_fast_transparent_png(self): + inference_options = [] + save_options = [] + + class FakeImage: + def __enter__(self): + return self + + def __exit__(self, *_args): + return None + + def convert(self, mode): + self.mode = mode + return self + + def save(self, path, **options): + save_options.append(options) + Path(path).touch() + + class FakeModel: + @classmethod + def from_pretrained(cls, model_dir): + self.assertEqual(model_dir, "model") + return cls() + + def to(model_self, device): + self.assertEqual(device.type, "cpu") + return model_self + + def eval(self): + return None + + def inference(self, image, **options): + inference_options.append((image, options)) + return FakeImage() + + fake_torch = ModuleType("torch") + fake_torch.cuda = SimpleNamespace( + is_available=lambda: False, + synchronize=lambda: None, + ) + fake_torch.device = lambda device_type: SimpleNamespace(type=device_type) + fake_torch.inference_mode = nullcontext + fake_ben2 = ModuleType("ben2") + fake_ben2.BEN_Base = FakeModel + fake_pil = ModuleType("PIL") + fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage()) + + with tempfile.TemporaryDirectory() as directory, patch.dict( + sys.modules, + { + "torch": fake_torch, + "ben2": fake_ben2, + "PIL": fake_pil, + }, + ): + root = Path(directory) + input_path = root / "input.png" + input_path.touch() + args = SimpleNamespace( + input=str(input_path), + model_dir="model", + output_dir=str(root / "result"), + status=str(root / "status.json"), + device="cpu", + ) + self.worker.remove_background(args) + status = json.loads(Path(args.status).read_text(encoding="utf-8")) + + self.assertTrue((root / "result" / "foreground.png").is_file()) + + self.assertEqual( + inference_options[0][1], + {"refine_foreground": False}, + ) + self.assertEqual( + save_options[0], + {"format": "PNG", "compress_level": 1}, + ) + self.assertIn("Background removed in", status["message"]) + def test_uses_official_da3_export_parameters_for_all_four_formats(self): inference_options = [] export_calls = [] diff --git a/uv.lock b/uv.lock index 8f5ad45..1f1d316 100644 --- a/uv.lock +++ b/uv.lock @@ -4,5 +4,5 @@ requires-python = ">=3.11" [[package]] name = "snapform" -version = "0.9.27" +version = "0.9.28" source = { virtual = "." }