From 07d836d9ad9f7d31512bbcea38798bbba5655bfb Mon Sep 17 00:00:00 2001 From: nan Date: Tue, 11 Aug 2026 13:23:13 +0800 Subject: [PATCH] Add geometry nodes asset and validation script - Implemented a new geometry nodes asset in `build_geometry_nodes_asset.py` to handle mesh refinement and topology fixing. - Created a validation script in `validate_geometry_nodes_asset.py` to ensure the integrity and functionality of the geometry node group. - Added input handling for various parameters including resolution, iterations, and smoothing options. - Included checks for node group structure and attributes to validate the geometry processing pipeline. --- README.md | 19 +- docs/technical.md | 57 ++- src/snapform/__init__.py | 5 + src/snapform/assets/O_SnapForm.blend | 3 + src/snapform/image_depth.py | 45 +- src/snapform/image_mesh.py | 564 ++++++++++++++++++++++ src/snapform/operators.py | 72 ++- src/snapform/runtime.py | 4 + src/snapform/worker.py | 468 ++++++++++++++++-- tests/test_addon.py | 178 +++++++ tests/test_image_depth.py | 2 +- tests/test_pack.py | 4 + tests/test_worker.py | 234 ++++++++- tools/build_geometry_nodes_asset.py | 633 +++++++++++++++++++++++++ tools/validate_geometry_nodes_asset.py | 188 ++++++++ 15 files changed, 2408 insertions(+), 68 deletions(-) create mode 100644 src/snapform/assets/O_SnapForm.blend create mode 100644 src/snapform/image_mesh.py create mode 100644 tools/build_geometry_nodes_asset.py create mode 100644 tools/validate_geometry_nodes_asset.py diff --git a/README.md b/README.md index 59f70d9..2c4366d 100644 --- a/README.md +++ b/README.md @@ -5,6 +5,7 @@ SnapForm 是一个 Blender 扩展，可以把一张照片、一段视频或一 ## 它能做什么 - 一键把 Image Empty 转换为立体面片：让 AI 估算深度，或直接用原图作为深度 +- 把单张 Image Empty 转换为透视纹理网格，并自动匹配原对象的 Transform 与 Offset - 一键去除图片背景，生成透明前景图 - 支持视频与序列帧输入，逐帧处理，动画也能得到逐帧对应的立体面片 - 全部操作都在右键菜单里完成，不需要写任何代码 @@ -27,7 +28,7 @@ DA3、BEN2 源码包和 uv 已内置在扩展里，安装时直接从本地解 ### 转换为立体面片 1. 在 3D 视图中选中一个 Image Empty（空物体，显示方式为 Image） -2. 右键选择 `Convert to Depth Plane` +2. 右键打开 `SnapForm` 子菜单，选择 `Convert to Depth Plane` 3. 在对话框中选择： - `Depth from`：`Original` 直接用原图，`AI Generated` 让 AI 估算深度 - `Background Depth Distance` 为 `0` 时不剔除背景；大于 `0` 时把背景设为固定距离，AI 与 Original 均支持 @@ -36,7 +37,21 @@ DA3、BEN2 源码包和 uv 已内置在扩展里，安装时直接从本地解 ### 去除背景 -选中 Image Empty 后，右键选择 `Remove Background`，完成后 Empty 会自动换成透明前景图。 +选中 Image Empty 后，右键打开 `SnapForm` 子菜单并选择 `Remove Background`，完成后 Empty 会自动 +换成透明前景图。该操作在子菜单中使用独立分割区域。 + +### 转换为透视网格 + +选中单张图片的 Image Empty，右键打开 `SnapForm` 子菜单并选择 `Convert to Mesh`。SnapForm 会按 DA3 +相机内参与深度逐像素反投影，并把生成网格匹配到原 Empty 的可见尺寸、Offset 与 Transform。DA3 +confidence、边拉伸与 BEN2 背景结果会保存为 Mesh attributes，再由一个可复用的 Geometry Nodes +修改器完成剔面、快速重拓扑、Relax 与 UV 重投。匹配相机默认不创建；开启后会作为 Mesh 的子级。 +当前此功能仅支持单张图片。 + +转换对话框只保留模型设置、`Create Matching Camera`、`Unshaded` 与 `Delete Original`。剔除和重拓扑参数直接在 +Geometry Nodes 修改器中调整；`Unshaded` 可使用与 Depth Plane 相同的 Emission 材质。普通材质的 +Bump Distance 会按透视网格的深度空间缩放。Mesh attributes 统一使用 `o_` 前缀；Host 准备过程显示 +`Post-processing n/6`，可看到当前阶段与总步骤。 ### 视频与序列帧 diff --git a/docs/technical.md b/docs/technical.md index 2d464db..c9a9f92 100644 --- a/docs/technical.md +++ b/docs/technical.md @@ -81,7 +81,7 @@ ## Image Empty 去背景 1. 确认 `Host > Models` 中的 `BEN2 Base` 已就绪；首次安装 Runtime 时会默认下载。 -2. 在 3D View 选中 Image Empty，右键选择 `Remove Background`。 +2. 在 3D View 选中 Image Empty，右键打开 `SnapForm` 子菜单并选择 `Remove Background`。 3. 完成后以 Fake User 保留原 Image datablock，新建并 Pack 一张透明 PNG 到当前 `.blend`，然后只把 当前 Empty 的 `data` 换成新图；Empty 的 Transform、显示尺寸、Offset 及其他设置保持不变。 4. 新 Image 使用原图去掉格式后再追加 `.foreground` 的名称。例如 `img.png` 生成 @@ -103,7 +103,7 @@ Cyclic 沿用原 Empty 的设置。序列帧无法被 Blender 整体 Pack，因 ## Image Empty 转深度面片 1. 在 3D View 选中 `Object Type = Empty`、`Display As = Image` 的对象。 -2. 右键选择 `Convert to Depth Plane`，通过 `Depth from` 选择深度来源： +2. 右键打开 `SnapForm` 子菜单并选择 `Convert to Depth Plane`，通过 `Depth from` 选择深度来源： - `AI Generated`：显示 Model、Use Ray Pose 和 Resolution Scale。实际 Process Res 为图片长边像素数 乘以 Resolution Scale；范围为 `0.1–1.0`，默认 `1.0`。计算后的 Process Res 会直接显示在 `Resolution Scale (… px)` 标签中。 @@ -144,6 +144,58 @@ Cyclic 沿用原 Empty 的设置。序列帧无法被 Blender 整体 Pack，因 `8×12` 个面。随后的 Simple Subdivision 修改器使用原生名称，Viewport Levels 固定为 `1`，Render Levels 固定为 `2`，不再随 Resolution Scale 或 Subdivision 动态变化。 +## Image Empty 转透视网格 + +1. 在 3D View 选中仅包含一帧的 Image Empty，右键打开 `SnapForm` 子菜单并选择 + `Convert to Mesh`。该入口独立于 + `Convert to Depth Plane`，当前不处理视频或序列帧。 +2. Host 请求 DA3 的 `depth`、`conf` 与 `intrinsics`，同时固定运行 BEN2 取得前景 alpha。Depth 不做 + smoothing，也不在 Host 中剔面或减面；Host 只反投影完整有效拓扑，并计算 confidence percentile、 + camera-space edge stretch 与 background 标记。 +3. Blender 把勾选 `Usage > Modifier` 的 `O Refine DA3 Mesh` 节点组直接作为 Geometry Nodes 修改器 + 添加到 Mesh，不另设修改器名称。节点组内没有子节点组，依次完成： + - 修改器顶层依次显示 `Resolution`、`Relax Iterations`、`Smooth`；默认展开的 `Culling` Panel + 包含 `Confidence`、`Stretch`、`Background`。`Stretch = 100%` 明确选择全部面，剔面后同时删除 + Face Count 为零的孤立顶点； + - 每轮 Remesh Iteration 依次执行 Connected Merge、循环拓扑修复与一次 Smooth，使新拓扑在进入下一轮 + Merge 前先稳定；拓扑修复删除非流形面、单面拓扑和小面积 mesh islands； + - 默认 `Resolution` 为原始有效三角边长中位数乘以 `sqrt(10) * 2`。`sqrt(10)` 来自原先 90% 面数缩减 + 的二维边距估计，再乘以 2 为迭代合并留出更平滑、稳定的默认尺度；节点内部 Connected Merge 的 + Distance 为 `Resolution * 0.5`； + - Smooth 采用 `O Smooth` 的边界检测和模糊法线权重，固定开放边界与薄面，只保留一个 + `Smooth` 程度系数；节点组不再额外执行 Flip Faces； + - 从修改前的原始表面以 nearest-surface 重新采样 `UVMap`，最终明确写为 Corner 域 `2D Vector`。 +4. 节点组随插件放在 `O_SnapForm.blend`。Blender 4.5 首次 Append 后复用同一 datablock； + Blender 5.0 以 linked ID 载入后调用 Pack，后续 Mesh 继续复用。插件不会为每个 Mesh 新建节点组， + 也不会在添加后重命名 Geometry Nodes 修改器。节点组只启用 `Usage > Modifier`，不 Mark as Asset； + 修改器关闭 `show_group_selector`。`Fix Topology` 位于默认折叠的 `Options` Panel。 + 节点图按 Culling、Remesh、Topology Fix、Island Cleanup 与 UV 分区摆放；各区就近使用独立的 + Group Input，并隐藏没有连线的接口，避免跨越整个节点图的输入连线。 +5. Host 准备过程固定显示 6 个编号阶段：验证 DA3 数据、反投影、准备 DA3/BEN2 attributes、构建原始 + 拓扑、投影 UV/attributes、写出网格。状态栏显示 `Post-processing n/6: ...`。 +6. Host 在 job 根目录写出 `mesh.npz` 与 `texture.png`。NPZ 保存 camera-space `vertices`、 + `faces`、`uv`、`intrinsics`、`image_size` 与有效深度中位数 `reference_depth`。Blender 读取这些 + 数据后直接创建 Mesh 与原图材质，不需要 Open3D 或 OpenCV 进入 Blender Python 环境。 + Mesh 保留下列结构化 attributes： + - 点域：`o_depth`、`o_confidence`、`o_confidence_percentile`、`o_background`； + - 面域：`o_face_confidence`、`o_edge_stretch`、`o_background_face`。 +7. 网格以 `reference_depth` 截面匹配 Image Empty 的显示矩形：该截面的投影四角与 Empty 的可见 + 边界重合，主点同时匹配 Empty Offset。网格对象直接沿用源 Empty 的 `matrix_world`，因此位置、旋转、 + 缩放及父级关系产生的世界变换保持一致。更近的顶点沿 Empty 局部 `+Z` 方向凸出，三角面正面朝向 + 观察相机，启用 Backface Culling 时可从正面看到。 +8. 转换对话框不再提供后处理参数，只保留模型设置、`Create Matching Camera`、`Unshaded` 与 + `Delete Original`。`Unshaded` 使用与 Depth Plane 相同的 Emission/Transparent 材质路径；普通材质的 + Bump Distance 按 `abs(object.scale.z) × camera scale_z` 缩放，与透视网格实际深度尺度一致。 + 相机默认关闭；开启后会在 Empty 局部 `+Z` 侧创建一台朝向局部 `-Z` 的透视相机，并成为生成 Mesh + 的子级，使二者移动时保持投影关系。开启 `Delete Original` 时，网格成功创建后删除源 Empty。 + 网格、材质与相机会归入同一个 `Convert to Mesh` Undo 步骤。 + +Debug 的 Textured Mesh GLB 会把 depth、confidence 与 confidence percentile 写入自定义顶点 +attributes。其他 DA3 Debug 三维输出会同时写出未压缩的 `prediction.attributes.npz`，保存完整 +depth、confidence、intrinsics、extrinsics 与 confidence percentile。当前 trimesh 的 glTF 点图元 +导出器只写 Position/Color，无法可靠把任意 attribute 嵌入点云 GLB，因此点云先以 companion NPZ +保留完整结构；后续若改为 SnapForm 自有 glTF writer，可再把这些字段直接嵌入点图元。 + 所有任务使用统一的状态栏进度显示。变化中的阶段说明显示在右侧的进度条内部，取消按钮 位于进度条右侧。状态字段支持阶段、当前操作与预计剩余秒数，界面统一组合为紧凑文案并保证 进度单调递增。任务进程与 Host 请求在后台启动，按钮点击后立即进入进度状态。Image Depth 的阶段包括 Runtime 启动、输入准备、模型加载、深度估算、缓存 @@ -234,6 +286,7 @@ Blender 为 Packed Image 或序列帧生成的输入副本也归档到对应 job - `download-model` - `export-reconstruction` - `generate-depth-plane` +- `generate-image-mesh` - `remove-background` - `remove-image-background` - `generate-original-depth-plane` diff --git a/src/snapform/__init__.py b/src/snapform/__init__.py index 4c81987..416b211 100644 --- a/src/snapform/__init__.py +++ b/src/snapform/__init__.py @@ -7,8 +7,10 @@ from .image_depth import ( BuildOriginalDepthPlane, ConvertImageToDepthPlane, RemoveImageBackground, + SnapFormImageMenu, draw_image_context_menu, ) +from .image_mesh import BuildImageMesh, ConvertImageToMesh from .operators import ( CancelJob, ControlHost, @@ -34,11 +36,14 @@ CLASSES = ( RunJob, CancelJob, ControlHost, + SnapFormImageMenu, ApplyRemovedBackground, RemoveImageBackground, BuildDepthPlane, BuildOriginalDepthPlane, ConvertImageToDepthPlane, + BuildImageMesh, + ConvertImageToMesh, HostPanel, ModelsPanel, DebugPanel, diff --git a/src/snapform/assets/O_SnapForm.blend b/src/snapform/assets/O_SnapForm.blend new file mode 100644 index 0000000..a7d1719 --- /dev/null +++ b/src/snapform/assets/O_SnapForm.blend @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:59a1b7a1b2b07b04cab23c5ea4609dffb912d49b411d0798bb8a0852cb179b67 +size 695763 diff --git a/src/snapform/image_depth.py b/src/snapform/image_depth.py index 5f7b421..31ad165 100644 --- a/src/snapform/image_depth.py +++ b/src/snapform/image_depth.py @@ -1783,22 +1783,37 @@ class ConvertImageToDepthPlane(bpy.types.Operator): return {"FINISHED"} +class SnapFormImageMenu(bpy.types.Menu): + bl_idname = "SNAPFORM_MT_image_context" + bl_label = "SnapForm" + + def draw(self, _context): + layout = self.layout + layout.operator_context = "INVOKE_DEFAULT" + layout.operator( + "snapform.convert_image_to_depth_plane", + text="Convert to Depth Plane", + icon="MOD_DISPLACE", + ) + layout.operator( + "snapform.convert_image_to_mesh", + text="Convert to Mesh", + icon="MESH_GRID", + ) + layout.separator() + layout.operator( + "snapform.remove_image_background", + text="Remove Background", + icon="IMAGE_ALPHA", + ) + + def draw_image_context_menu(self, context): - active_object = context.object - if not is_image_empty(active_object): + if not is_image_empty(context.object): return self.layout.separator() - row = self.layout.row() - row.operator_context = "INVOKE_DEFAULT" - row.operator( - "snapform.convert_image_to_depth_plane", - 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", + self.layout.menu( + SnapFormImageMenu.bl_idname, + text="SnapForm", + icon="PLUGIN", ) diff --git a/src/snapform/image_mesh.py b/src/snapform/image_mesh.py new file mode 100644 index 0000000..44be6ed --- /dev/null +++ b/src/snapform/image_mesh.py @@ -0,0 +1,564 @@ +from pathlib import Path + +import bpy + +from .image_depth import ( + create_material, + image_depth_model_items, + is_image_empty, + plane_bounds, + prepare_image_input, + remove_temporary_input, + scaled_process_resolution, + video_total_frames, +) + + +_last_mesh_model = "BASE" + +POINT_ATTRIBUTES = { + "o_depth": ("FLOAT", "point_depth"), + "o_confidence": ("FLOAT", "point_confidence"), + "o_confidence_percentile": ( + "FLOAT", + "point_confidence_percentile", + ), + "o_background": ("BOOLEAN", "point_background"), +} +FACE_ATTRIBUTES = { + "o_face_confidence": ( + "FLOAT", + "face_confidence_percentile", + ), + "o_edge_stretch": ("FLOAT", "face_edge_stretch"), + "o_background_face": ("BOOLEAN", "face_background"), +} +NODE_GROUP_NAME = "O Refine DA3 Mesh" +NODE_ASSET_NAME = "O_SnapForm.blend" + + +def load_image_mesh_data(npz_path): + import numpy as np + + try: + with np.load(npz_path, allow_pickle=False) as archive: + required = ( + "vertices", + "faces", + "uv", + "intrinsics", + "image_size", + "reference_depth", + "point_depth", + "point_confidence", + "point_confidence_percentile", + "point_background", + "face_confidence_percentile", + "face_edge_stretch", + "face_background", + ) + missing = [key for key in required if key not in archive.files] + if missing: + raise ValueError( + f"Mesh data is missing: {', '.join(missing)}" + ) + vertices = np.asarray(archive["vertices"], dtype=np.float32) + faces = np.asarray(archive["faces"], dtype=np.int32) + uv = np.asarray(archive["uv"], dtype=np.float32) + intrinsics = np.asarray( + archive["intrinsics"], + dtype=np.float32, + ) + image_size = tuple( + int(value) for value in archive["image_size"].tolist() + ) + reference_depth = float(archive["reference_depth"]) + attributes = { + key: np.asarray(archive[key]) + for _name, (_data_type, key) in { + **POINT_ATTRIBUTES, + **FACE_ATTRIBUTES, + }.items() + } + except OSError as error: + raise ValueError(f"Unable to read mesh data: {npz_path}") from error + if vertices.ndim != 2 or vertices.shape[1] != 3: + raise ValueError(f"Unsupported vertices shape: {vertices.shape}") + if faces.ndim != 2 or faces.shape[1] != 3: + raise ValueError(f"Unsupported faces shape: {faces.shape}") + if uv.shape != (len(vertices), 2): + raise ValueError("Mesh UV and vertex counts do not match") + if intrinsics.shape != (3, 3): + raise ValueError(f"Unsupported intrinsics shape: {intrinsics.shape}") + if len(image_size) != 2 or min(image_size) <= 0: + raise ValueError(f"Unsupported image size: {image_size}") + if not reference_depth > 0.0: + raise ValueError("Mesh reference depth must be positive") + for _name, (_data_type, key) in POINT_ATTRIBUTES.items(): + if len(attributes[key]) != len(vertices): + raise ValueError(f"Point attribute has the wrong length: {key}") + for _name, (_data_type, key) in FACE_ATTRIBUTES.items(): + if len(attributes[key]) != len(faces): + raise ValueError(f"Face attribute has the wrong length: {key}") + return ( + vertices, + faces, + uv, + intrinsics, + image_size, + reference_depth, + attributes, + ) + + +def empty_anchored_vertices( + source_object, + camera_vertices, + intrinsics, + image_size, + reference_depth, +): + import numpy as np + + width, height = image_size + x_min, x_max, y_min, y_max = plane_bounds(source_object) + plane_width = x_max - x_min + plane_height = y_max - y_min + pixel_width = max(width - 1, 1) + pixel_height = max(height - 1, 1) + fx = float(intrinsics[0, 0]) + fy = float(intrinsics[1, 1]) + cx = float(intrinsics[0, 2]) + cy = float(intrinsics[1, 2]) + if fx <= 0.0 or fy <= 0.0: + raise ValueError("Camera intrinsics contain an invalid focal length") + scale_x = plane_width * fx / (pixel_width * reference_depth) + scale_y = plane_height * fy / (pixel_height * reference_depth) + scale_z = float(np.sqrt(scale_x * scale_y)) + camera_x = x_min + plane_width * cx / pixel_width + camera_y = y_max - plane_height * cy / pixel_height + vertices = np.empty_like(camera_vertices, dtype=np.float32) + vertices[:, 0] = camera_x + camera_vertices[:, 0] * scale_x + vertices[:, 1] = camera_y - camera_vertices[:, 1] * scale_y + vertices[:, 2] = ( + reference_depth - camera_vertices[:, 2] + ) * scale_z + camera_location = ( + camera_x, + camera_y, + reference_depth * scale_z, + ) + camera_parameters = { + "scale_x": scale_x, + "scale_y": scale_y, + "scale_z": scale_z, + "location": camera_location, + } + return vertices, camera_parameters + + +def assign_vertex_uv(mesh, vertex_uv): + import numpy as np + + uv_layer = mesh.uv_layers.new(name="UVMap") + loop_vertices = np.empty(len(mesh.loops), dtype=np.int32) + mesh.loops.foreach_get("vertex_index", loop_vertices) + loop_uv = vertex_uv[loop_vertices] + uv_layer.data.foreach_set("uv", loop_uv.reshape(-1)) + + +def assign_mesh_attribute(mesh, name, data_type, domain, values): + import numpy as np + + attribute = mesh.attributes.new( + name=name, + type=data_type, + domain=domain, + ) + array = np.asarray(values) + if data_type == "BOOLEAN": + array = array.astype(bool) + else: + array = array.astype(np.float32) + attribute.data.foreach_set("value", array.reshape(-1)) + + +def assign_snapform_attributes(mesh, attributes): + for name, (data_type, key) in POINT_ATTRIBUTES.items(): + assign_mesh_attribute( + mesh, + name, + data_type, + "POINT", + attributes[key], + ) + for name, (data_type, key) in FACE_ATTRIBUTES.items(): + assign_mesh_attribute( + mesh, + name, + data_type, + "FACE", + attributes[key], + ) + + +def node_asset_path(): + return Path(__file__).with_name("assets") / NODE_ASSET_NAME + + +def load_geometry_node_group(): + existing = bpy.data.node_groups.get(NODE_GROUP_NAME) + if existing is not None: + return existing + asset_path = node_asset_path() + if not asset_path.is_file(): + raise RuntimeError(f"Missing Geometry Nodes asset: {asset_path}") + use_pack = bpy.app.version >= (5, 0, 0) + if not use_pack: + directory = str(asset_path) + "/NodeTree/" + result = bpy.ops.wm.append( + filepath=directory + NODE_GROUP_NAME, + directory=directory, + filename=NODE_GROUP_NAME, + link=False, + do_reuse_local_id=True, + ) + if "FINISHED" not in result: + raise RuntimeError("Unable to append the Geometry Nodes asset") + node_group = bpy.data.node_groups.get(NODE_GROUP_NAME) + if node_group is None: + raise RuntimeError("Appended Geometry Nodes asset is unavailable") + node_group.use_fake_user = True + return node_group + with bpy.data.libraries.load(str(asset_path), link=use_pack) as ( + source, + target, + ): + if NODE_GROUP_NAME not in source.node_groups: + raise RuntimeError( + f"Geometry Nodes asset is missing {NODE_GROUP_NAME}" + ) + target.node_groups = [NODE_GROUP_NAME] + node_group = target.node_groups[0] + if use_pack: + linked_group = node_group + node_group = bpy.data.pack_linked_ids_hierarchy(root_id=linked_group) + if linked_group != node_group and linked_group.users == 0: + bpy.data.node_groups.remove(linked_group) + node_group.use_fake_user = True + return node_group + + +def modifier_input_identifier(node_group, name): + for item in node_group.interface.items_tree: + if item.item_type == "SOCKET" and item.in_out == "INPUT": + if item.name == name: + return item.identifier + raise RuntimeError(f"Geometry Nodes input is missing: {name}") + + +def default_remesh_resolution(vertices, faces): + import numpy as np + + if not len(faces): + return 0.01 + triangles = np.asarray(vertices, dtype=np.float32)[faces] + edge_lengths = np.linalg.norm( + triangles - np.roll(triangles, 1, axis=1), + axis=2, + ) + valid = edge_lengths[np.isfinite(edge_lengths) & (edge_lengths > 0.0)] + if not len(valid): + return 0.01 + # Twice the 90% reduction estimate gives the iterative remesh more room + # to settle into a smooth, stable default topology. + return float(np.median(valid) * np.sqrt(10.0) * 2.0) + + +def material_depth_scale(mesh_object, camera_parameters): + return abs(float(mesh_object.scale.z)) * abs( + float(camera_parameters["scale_z"]) + ) + + +def add_geometry_nodes_processing(mesh_object, vertices, faces): + node_group = load_geometry_node_group() + modifier = mesh_object.modifiers.new(name=node_group.name, type="NODES") + modifier.node_group = node_group + modifier.show_group_selector = False + identifier = modifier_input_identifier(node_group, "Resolution") + modifier[identifier] = default_remesh_resolution(vertices, faces) + + +def create_matching_camera( + context, + source_object, + mesh_object, + name, + intrinsics, + image_size, + camera_parameters, +): + from mathutils import Matrix + + width, height = image_size + scale_x = camera_parameters["scale_x"] + scale_z = camera_parameters["scale_z"] + effective_fx = float(intrinsics[0, 0]) * scale_z / scale_x + camera_data = bpy.data.cameras.new(f"{name}.camera") + camera_data.sensor_fit = "HORIZONTAL" + camera_data.sensor_width = 36.0 + camera_data.lens = effective_fx * camera_data.sensor_width / max( + width - 1, + 1, + ) + cx = float(intrinsics[0, 2]) + cy = float(intrinsics[1, 2]) + camera_data.shift_x = (cx - (width - 1) * 0.5) / max(width - 1, 1) + camera_data.shift_y = ((height - 1) * 0.5 - cy) / max(width - 1, 1) + camera_object = bpy.data.objects.new(f"{name}.camera", camera_data) + context.collection.objects.link(camera_object) + local_matrix = Matrix.Translation(camera_parameters["location"]) + camera_object.matrix_world = source_object.matrix_world @ local_matrix + camera_world = camera_object.matrix_world.copy() + camera_object.parent = mesh_object + camera_object.matrix_world = camera_world + return camera_object + + +def build_image_mesh( + context, + source_object, + npz_path, + delete_source, + create_camera, + unshaded, +): + ( + camera_vertices, + faces, + uv, + intrinsics, + image_size, + reference_depth, + attributes, + ) = load_image_mesh_data(npz_path) + vertices, camera_parameters = empty_anchored_vertices( + source_object, + camera_vertices, + intrinsics, + image_size, + reference_depth, + ) + source_name = source_object.name + mesh = bpy.data.meshes.new(source_name) + mesh.from_pydata(vertices.tolist(), [], faces.tolist()) + assign_vertex_uv(mesh, uv) + assign_snapform_attributes(mesh, attributes) + mesh.update() + for polygon in mesh.polygons: + polygon.use_smooth = True + mesh_object = bpy.data.objects.new(source_name, mesh) + context.collection.objects.link(mesh_object) + mesh_object.matrix_world = source_object.matrix_world.copy() + depth_scale = material_depth_scale(mesh_object, camera_parameters) + material = create_material( + source_name, + source_object.data, + depth_scale, + unshaded=unshaded, + ) + mesh.materials.append(material) + mesh_object["snapform_reference_depth"] = reference_depth + mesh_object["snapform_source_type"] = "IMAGE_MESH" + add_geometry_nodes_processing(mesh_object, vertices, faces) + if create_camera: + create_matching_camera( + context, + source_object, + mesh_object, + source_name, + intrinsics, + image_size, + camera_parameters, + ) + for selected in context.selected_objects: + selected.select_set(False) + mesh_object.select_set(True) + context.view_layer.objects.active = mesh_object + if delete_source: + bpy.data.objects.remove(source_object, do_unlink=True) + mesh_object.name = source_name + return mesh_object + + +class BuildImageMesh(bpy.types.Operator): + bl_idname = "snapform.build_image_mesh" + bl_label = "Build Image Mesh" + bl_options = {"INTERNAL"} + + snapform_source_object_name: bpy.props.StringProperty( + options={"HIDDEN", "SKIP_SAVE"}, + ) + snapform_npz_path: bpy.props.StringProperty( + options={"HIDDEN", "SKIP_SAVE"}, + ) + snapform_delete_source: bpy.props.BoolProperty( + default=True, + options={"HIDDEN", "SKIP_SAVE"}, + ) + snapform_create_camera: bpy.props.BoolProperty( + default=False, + 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 no longer exists") + return {"CANCELLED"} + try: + build_image_mesh( + context, + source_object, + self.snapform_npz_path, + self.snapform_delete_source, + self.snapform_create_camera, + self.snapform_unshaded, + ) + except (OSError, RuntimeError, ValueError) as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + return {"FINISHED"} + + +class ConvertImageToMesh(bpy.types.Operator): + bl_idname = "snapform.convert_image_to_mesh" + bl_label = "Convert to Mesh" + bl_description = "Convert an Image Empty into a perspective textured mesh" + bl_options = {"REGISTER"} + + snapform_model: bpy.props.EnumProperty( + name="Model", + description="Model used to estimate image depth", + items=image_depth_model_items, + ) + snapform_use_ray_pose: bpy.props.BoolProperty( + name="Use Ray Pose", + description="Use ray-based camera pose estimation", + default=True, + ) + snapform_process_res_factor: bpy.props.FloatProperty( + name="Resolution Scale", + description="Scale applied to the source image long-edge resolution", + default=1.0, + min=0.1, + max=1.0, + step=10, + precision=2, + ) + snapform_create_camera: bpy.props.BoolProperty( + name="Create Matching Camera", + description="Create a perspective camera matching the source projection", + default=False, + ) + snapform_unshaded: bpy.props.BoolProperty( + name="Unshaded", + description="Display the source color through an Emission shader", + default=False, + ) + snapform_delete_original: bpy.props.BoolProperty( + name="Delete Original", + description="Delete the source Image Empty after conversion", + default=True, + ) + + @classmethod + def poll(cls, context): + from .models import BEN2_MODEL, MODELS + from .operators import active_job + from .runtime import host_busy, model_ready + + return ( + is_image_empty(context.object) + and video_total_frames(context.object.data) == 1 + and active_job() is None + and not host_busy() + and any(model_ready(model_key) for model_key in MODELS) + and model_ready(BEN2_MODEL.key) + ) + + def invoke(self, context, _event): + global _last_mesh_model + + if not self.properties.is_property_set("snapform_model"): + available_keys = [ + item[0] for item in image_depth_model_items(self, context) + ] + if available_keys: + self.snapform_model = ( + _last_mesh_model + if _last_mesh_model in available_keys + else available_keys[0] + ) + return context.window_manager.invoke_props_dialog(self, width=360) + + def draw(self, context): + layout = self.layout + layout.prop(self, "snapform_model") + layout.prop(self, "snapform_use_ray_pose") + process_res = scaled_process_resolution( + context.object.data, + self.snapform_process_res_factor, + ) + layout.prop( + self, + "snapform_process_res_factor", + text=f"Resolution Scale ({process_res} px)", + ) + layout.prop(self, "snapform_create_camera") + layout.prop(self, "snapform_unshaded") + layout.prop(self, "snapform_delete_original") + + def execute(self, context): + from .operators import IMAGE_MESH_JOB + + source_object = context.object + if video_total_frames(source_object.data) != 1: + self.report({"ERROR"}, "Convert to Mesh currently supports still images") + return {"CANCELLED"} + try: + process_res = scaled_process_resolution( + source_object.data, + self.snapform_process_res_factor, + ) + source_path, delete_input = prepare_image_input( + source_object.data, + context.scene, + ) + except (RuntimeError, ValueError) as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + result = bpy.ops.snapform.run_job( + job_type=IMAGE_MESH_JOB, + snapform_input_path=str(source_path), + snapform_source_object_name=source_object.name, + snapform_model_key=self.snapform_model, + snapform_use_ray_pose=self.snapform_use_ray_pose, + snapform_process_res=process_res, + snapform_delete_source=self.snapform_delete_original, + snapform_delete_input=delete_input, + snapform_mesh_create_camera=self.snapform_create_camera, + snapform_unshaded=self.snapform_unshaded, + ) + if "CANCELLED" in result: + if delete_input: + remove_temporary_input(source_path) + return {"CANCELLED"} + global _last_mesh_model + _last_mesh_model = self.snapform_model + return {"FINISHED"} diff --git a/src/snapform/operators.py b/src/snapform/operators.py index 9424735..c276c9c 100644 --- a/src/snapform/operators.py +++ b/src/snapform/operators.py @@ -17,6 +17,7 @@ from .models import ( RECONSTRUCTION_JOB = "RECONSTRUCT" IMAGE_DEPTH_JOB = "IMAGE_DEPTH" +IMAGE_MESH_JOB = "IMAGE_MESH" BACKGROUND_REMOVAL_JOB = "REMOVE_BACKGROUND" IMAGE_BACKGROUND_JOB = "IMAGE_BACKGROUND" ORIGINAL_BACKGROUND_JOB = "ORIGINAL_BACKGROUND" @@ -78,6 +79,7 @@ class ActiveJob: unshaded: bool = False remove_background: bool = False background_depth_distance: float = 0.0 + mesh_create_camera: bool = False temporary_input_path: Path | None = None cancelled: bool = False uses_host: bool = False @@ -139,6 +141,7 @@ class RunJob(bpy.types.Operator): ("DOWNLOAD_MODEL", "Download Model", ""), (RECONSTRUCTION_JOB, "Export DA3 Result", ""), (IMAGE_DEPTH_JOB, "Convert Image to Depth Plane", ""), + (IMAGE_MESH_JOB, "Convert Image to Mesh", ""), (BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""), (IMAGE_BACKGROUND_JOB, "Remove Image Empty Background", ""), (ORIGINAL_BACKGROUND_JOB, "Build Original Background Plane", ""), @@ -203,7 +206,14 @@ class RunJob(bpy.types.Operator): min=0.0, options={"HIDDEN", "SKIP_SAVE"}, ) - + snapform_mesh_create_camera: 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 @@ -227,6 +237,8 @@ class RunJob(bpy.types.Operator): export_format = settings.export_format elif self.job_type == IMAGE_DEPTH_JOB: export_format = "mini_npz" + elif self.job_type == IMAGE_MESH_JOB: + export_format = "mesh_data" job = ActiveJob( kind=self.job_type, @@ -243,12 +255,21 @@ 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_background_depth_distance > 0.0, + remove_background=( + self.snapform_background_depth_distance > 0.0 + or bool(getattr(self, "snapform_remove_background", False)) + ), background_depth_distance=self.snapform_background_depth_distance, + mesh_create_camera=getattr( + self, + "snapform_mesh_create_camera", + False, + ), temporary_input_path=( Path(self.snapform_input_path) if self.job_type in { IMAGE_DEPTH_JOB, + IMAGE_MESH_JOB, IMAGE_BACKGROUND_JOB, ORIGINAL_BACKGROUND_JOB, ORIGINAL_FRAMES_JOB, @@ -317,6 +338,7 @@ class RunJob(bpy.types.Operator): getattr(self, "snapform_delete_input", False) and self.job_type in { IMAGE_DEPTH_JOB, + IMAGE_MESH_JOB, IMAGE_BACKGROUND_JOB, ORIGINAL_BACKGROUND_JOB, ORIGINAL_FRAMES_JOB, @@ -332,6 +354,8 @@ class RunJob(bpy.types.Operator): getattr(self, "snapform_background_depth_distance", 0.0) ) remove_background = background_depth_distance > 0.0 + if self.job_type == IMAGE_MESH_JOB: + remove_background = True if self.job_type == "INSTALL_RUNTIME": self._require_online_access() return runtime.install_runtime_command(paths["status"]), paths @@ -349,6 +373,7 @@ class RunJob(bpy.types.Operator): if self.job_type in { IMAGE_DEPTH_JOB, + IMAGE_MESH_JOB, IMAGE_BACKGROUND_JOB, ORIGINAL_BACKGROUND_JOB, ORIGINAL_FRAMES_JOB, @@ -410,15 +435,18 @@ class RunJob(bpy.types.Operator): if self.job_type not in { RECONSTRUCTION_JOB, IMAGE_DEPTH_JOB, + IMAGE_MESH_JOB, }: raise RuntimeError(f"Unsupported task type: {self.job_type}") export_format = settings.export_format if self.job_type == IMAGE_DEPTH_JOB: export_format = "mini_npz" + elif self.job_type == IMAGE_MESH_JOB: + export_format = "mesh_data" use_ray_pose = None process_res = None model_key = selected_model_key(settings, export_format) - if self.job_type == IMAGE_DEPTH_JOB: + if self.job_type in {IMAGE_DEPTH_JOB, IMAGE_MESH_JOB}: model_key = self.snapform_model_key use_ray_pose = self.snapform_use_ray_pose process_res = self.snapform_process_res @@ -426,7 +454,7 @@ class RunJob(bpy.types.Operator): if not runtime.model_ready(model_key): raise RuntimeError(f"Download {model.label} first") if ( - self.job_type == IMAGE_DEPTH_JOB + self.job_type in {IMAGE_DEPTH_JOB, IMAGE_MESH_JOB} and remove_background and not runtime.model_ready(BEN2_MODEL.key) ): @@ -446,13 +474,13 @@ class RunJob(bpy.types.Operator): process_res=process_res, process_res_method=( IMAGE_DEPTH_PROCESS_RES_METHOD - if self.job_type == IMAGE_DEPTH_JOB + if self.job_type in {IMAGE_DEPTH_JOB, IMAGE_MESH_JOB} else None ), use_ray_pose=use_ray_pose, start_frame=start_frame, remove_background=( - self.job_type == IMAGE_DEPTH_JOB + self.job_type in {IMAGE_DEPTH_JOB, IMAGE_MESH_JOB} and remove_background ), background_depth_distance=( @@ -469,6 +497,7 @@ class RunJob(bpy.types.Operator): """Debug panels use the scene device; production tasks use Preferences.""" if self.job_type in { IMAGE_DEPTH_JOB, + IMAGE_MESH_JOB, IMAGE_BACKGROUND_JOB, ORIGINAL_BACKGROUND_JOB, }: @@ -490,6 +519,8 @@ 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 == IMAGE_MESH_JOB: + return "Starting perspective mesh generation" if self.job_type == BACKGROUND_REMOVAL_JOB: return "Starting BEN2 background removal" if self.job_type == IMAGE_BACKGROUND_JOB: @@ -663,6 +694,21 @@ class RunJob(bpy.types.Operator): self._redraw_ui(context, force=True) return {"CANCELLED"} status["message"] = "Depth plane is ready" + if job.kind == IMAGE_MESH_JOB: + try: + settings.status = "Building perspective mesh" + settings.progress = 0.96 + self._redraw_ui(context, force=True) + self._build_image_mesh(job) + except (OSError, RuntimeError, ValueError) as import_error: + settings.status = "Mesh 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"} + status["message"] = "Perspective mesh is ready" if job.kind in {ORIGINAL_BACKGROUND_JOB, ORIGINAL_FRAMES_JOB}: try: settings.status = "Building depth plane" @@ -706,6 +752,20 @@ class RunJob(bpy.types.Operator): self._redraw_ui(context, force=True) return {"FINISHED"} + def _build_image_mesh(self, job): + result = bpy.ops.snapform.build_image_mesh( + snapform_source_object_name=job.source_object_name, + snapform_npz_path=str(job.result_directory / "mesh.npz"), + snapform_delete_source=job.delete_source, + snapform_create_camera=job.mesh_create_camera, + snapform_unshaded=job.unshaded, + ) + if "FINISHED" not in result: + raise RuntimeError("Unable to build the perspective mesh") + undo_result = bpy.ops.ed.undo_push(message="Convert to Mesh") + if "FINISHED" not in undo_result: + raise RuntimeError("Unable to create the mesh Undo step") + def _build_depth_plane(self, context, job): result_directory = job.result_directory npz_path = result_directory / "results.npz" diff --git a/src/snapform/runtime.py b/src/snapform/runtime.py index 9ab383a..4aca9cf 100644 --- a/src/snapform/runtime.py +++ b/src/snapform/runtime.py @@ -37,6 +37,7 @@ JOB_FOLDER_NAMES = { "DOWNLOAD_MODEL": "download-model", "RECONSTRUCT": "export-reconstruction", "IMAGE_DEPTH": "generate-depth-plane", + "IMAGE_MESH": "generate-image-mesh", "REMOVE_BACKGROUND": "remove-background", "IMAGE_BACKGROUND": "remove-image-background", "ORIGINAL_BACKGROUND": "generate-original-depth-plane", @@ -812,6 +813,7 @@ def reconstruct_request( use_ray_pose=None, remove_background=False, background_depth_distance=0.0, + mesh_filter_edges=None, start_frame=1, device=None, ): @@ -852,6 +854,8 @@ def reconstruct_request( "background_depth_distance": float(background_depth_distance), "mesh_filter_edges": bool( getattr(settings, "mesh_filter_edges", True) + if mesh_filter_edges is None + else mesh_filter_edges ), "mesh_max_edge_length": float( getattr(settings, "mesh_max_edge_length", 0.1) diff --git a/src/snapform/worker.py b/src/snapform/worker.py index 2fd04f3..7f95c22 100644 --- a/src/snapform/worker.py +++ b/src/snapform/worker.py @@ -24,7 +24,14 @@ VIDEO_SUFFIXES = { ".mpeg", } SCENE_FILENAME = "scene.glb" -EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz", "mesh") +EXPORT_FORMATS = ( + "glb", + "gs_ply", + "depth_vis", + "mini_npz", + "mesh", + "mesh_data", +) HF_ENDPOINT = "https://huggingface.co" MODELSCOPE_ENDPOINT = "https://modelscope.cn" MODELSCOPE_REVISION = "master" @@ -368,41 +375,357 @@ def unproject_depth_frame(depth, intrinsics, extrinsics): return world_points[:, :3].reshape(height, width, 3) +def unproject_camera_depth(depth, intrinsics): + import numpy as np + + height, width = depth.shape + x, y = np.meshgrid( + np.arange(width, dtype=np.float32), + np.arange(height, dtype=np.float32), + ) + pixels = np.stack((x, y, np.ones_like(x)), axis=-1).reshape(-1, 3) + rays = (np.linalg.inv(intrinsics) @ pixels.T).T + return (rays * depth.reshape(-1, 1)).reshape(height, width, 3) + + +def smooth_depth_map(depth, smoothing): + import numpy as np + + factor = min(max(float(smoothing) / 100.0, 0.0), 1.0) + if factor <= 0.0: + return np.asarray(depth, dtype=np.float32).copy() + import cv2 + + values = np.asarray(depth, dtype=np.float32) + valid = np.isfinite(values) & (values > 0.0) + if not valid.any(): + raise ValueError("DA3 returned no valid depth values") + low, high = np.percentile(values[valid], (5.0, 95.0)) + depth_range = max(float(high - low), 1e-6) + diameter = 3 + 2 * int(round(3.0 * factor)) + working_values = values.copy() + working_values[~valid] = float(np.median(values[valid])) + smoothed = cv2.bilateralFilter( + working_values, + diameter, + depth_range * (0.01 + 0.09 * factor), + 1.0 + 4.0 * factor, + ) + smoothed[~valid] = values[~valid] + return smoothed + + +def compact_mesh(vertices, faces): + import numpy as np + + vertices = np.asarray(vertices, dtype=np.float32) + faces = np.asarray(faces, dtype=np.int32) + used_vertices = np.unique(faces) + index_map = np.full(len(vertices), -1, dtype=np.int32) + index_map[used_vertices] = np.arange(len(used_vertices), dtype=np.int32) + vertices = vertices[used_vertices] + faces = index_map[faces] + return vertices, faces + + +def project_camera_uv(vertices, intrinsics, image_size): + import numpy as np + + width, height = image_size + z = vertices[:, 2] + if np.any(~np.isfinite(z)) or np.any(z <= 0.0): + raise ValueError("Simplified mesh contains invalid camera depth") + u = intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2] + v = intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2] + denominator_x = max(width - 1, 1) + denominator_y = max(height - 1, 1) + return np.column_stack( + ( + np.clip(u / denominator_x, 0.0, 1.0), + np.clip(1.0 - v / denominator_y, 0.0, 1.0), + ) + ).astype(np.float32) + + +def project_camera_pixels(vertices, intrinsics): + import numpy as np + + z = vertices[:, 2] + if np.any(~np.isfinite(z)) or np.any(z <= 0.0): + raise ValueError("Mesh contains invalid camera depth") + return np.column_stack( + ( + intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2], + intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2], + ) + ) + + +def sample_image_nearest(values, pixels): + import numpy as np + + height, width = values.shape + x = np.clip(np.rint(pixels[:, 0]).astype(np.int32), 0, width - 1) + y = np.clip(np.rint(pixels[:, 1]).astype(np.int32), 0, height - 1) + return np.asarray(values)[y, x] + + +def confidence_percentile_map(depth, confidence): + import numpy as np + + valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence) + result = np.zeros(depth.shape, dtype=np.float32) + if not valid.any(): + return result + sorted_values = np.sort(np.asarray(confidence)[valid]) + ranks = np.searchsorted( + sorted_values, + np.asarray(confidence)[valid], + side="right", + ) + result[valid] = ranks.astype(np.float32) / len(sorted_values) + return result + + +def triangle_edge_stretch_values(points, triangles): + import numpy as np + + flat_points = np.asarray(points, dtype=np.float32).reshape(-1, 3) + triangle_points = flat_points[triangles] + triangle_depth = triangle_points[..., 2] + safe_depth = np.where( + np.abs(triangle_depth) > 1e-8, + triangle_depth, + np.nan, + ) + triangle_rays = triangle_points / safe_depth[..., None] + next_points = np.roll(triangle_points, -1, axis=1) + next_rays = np.roll(triangle_rays, -1, axis=1) + next_depth = np.roll(triangle_depth, -1, axis=1) + local_depth = (triangle_depth + next_depth) * 0.5 + flat_edges = ( + triangle_rays * local_depth[..., None] + - next_rays * local_depth[..., None] + ) + flat_lengths = np.linalg.norm(flat_edges, axis=-1) + actual_lengths = np.linalg.norm( + triangle_points - next_points, + axis=-1, + ) + stretch = actual_lengths / np.maximum(flat_lengths, 1e-8) + stretch[np.isclose(stretch, 1.0, rtol=1e-5, atol=1e-6)] = 1.0 + return stretch + + +def image_mesh_attributes( + vertices, + faces, + intrinsics, + depth, + confidence, + confidence_percentiles, + foreground_mask, +): + import numpy as np + + pixels = project_camera_pixels(vertices, intrinsics) + point_confidence = sample_image_nearest(confidence, pixels).astype( + np.float32 + ) + point_confidence = np.nan_to_num(point_confidence, nan=0.0) + point_percentile = sample_image_nearest( + confidence_percentiles, + pixels, + ).astype(np.float32) + point_background = ~sample_image_nearest(foreground_mask, pixels).astype( + bool + ) + face_confidence = point_percentile[faces].min(axis=1) + face_edge_stretch = triangle_edge_stretch_values( + vertices, + faces, + ).max(axis=1) + face_background = point_background[faces].any(axis=1) + return { + "point_depth": vertices[:, 2].astype(np.float32), + "point_confidence": point_confidence, + "point_confidence_percentile": point_percentile, + "point_background": point_background, + "face_confidence_percentile": face_confidence.astype(np.float32), + "face_edge_stretch": face_edge_stretch.astype(np.float32), + "face_background": face_background, + } + + +def export_prediction_attributes(prediction, output_directory): + import numpy as np + + values = {} + for name in ("depth", "conf", "intrinsics", "extrinsics"): + value = getattr(prediction, name, None) + if value is not None: + values[name] = np.asarray(value) + if "depth" in values and "conf" in values: + values["confidence_percentile"] = np.stack( + [ + confidence_percentile_map(depth, confidence) + for depth, confidence in zip( + values["depth"], + values["conf"], + ) + ] + ) + output_path = Path(output_directory) / "prediction.attributes.npz" + np.savez(output_path, **values) + return output_path + + +def export_image_mesh_data( + prediction, + image_path, + output_directory, + foreground_path=None, + progress_callback=None, +): + import numpy as np + from PIL import Image + + total_steps = 6 + + def progress(step, label): + if progress_callback is not None: + progress_callback(step, total_steps, label) + + progress(1, "Validating DA3 data") + required = ("depth", "conf", "intrinsics") + missing = [ + name + for name in required + if getattr(prediction, name, None) is None + ] + if missing: + raise ValueError(f"DA3 mesh output is missing: {', '.join(missing)}") + depth_frames = np.asarray(prediction.depth, dtype=np.float32) + confidence_frames = np.asarray(prediction.conf, dtype=np.float32) + if depth_frames.shape[0] != 1: + raise ValueError("Convert to Mesh currently supports one image") + raw_depth = depth_frames[0] + confidence = confidence_frames[0] + intrinsics = np.asarray(prediction.intrinsics[0], dtype=np.float32) + progress(2, "Unprojecting camera geometry") + points = unproject_camera_depth(raw_depth, intrinsics) + progress(3, "Preparing DA3 and BEN2 attributes") + foreground_mask = np.ones(raw_depth.shape, dtype=bool) + if foreground_path is not None: + foreground_mask = foreground_mesh_mask( + foreground_path, + raw_depth.shape, + ) + confidence_percentiles = confidence_percentile_map( + raw_depth, + confidence, + ) + progress(4, "Building editable source topology") + faces = mesh_triangle_faces( + points, + raw_depth, + confidence, + None, + False, + 0.0, + ) + if not len(faces): + raise ValueError( + "DA3 output does not contain any valid mesh faces" + ) + progress(5, "Projecting UV and structured attributes") + vertices, faces = compact_mesh( + points.reshape(-1, 3), + faces, + ) + height, width = raw_depth.shape + uv = project_camera_uv(vertices, intrinsics, (width, height)) + attributes = image_mesh_attributes( + vertices, + faces, + intrinsics, + raw_depth, + confidence, + confidence_percentiles, + foreground_mask, + ) + valid_depth = raw_depth[np.isfinite(raw_depth) & (raw_depth > 0.0)] + reference_depth = float(np.median(valid_depth)) + output_directory = Path(output_directory) + output_directory.mkdir(parents=True, exist_ok=True) + progress(6, "Writing mesh data") + result_path = output_directory / "mesh.npz" + np.savez_compressed( + result_path, + vertices=vertices, + faces=faces, + uv=uv, + intrinsics=intrinsics, + image_size=np.asarray((width, height), dtype=np.int32), + reference_depth=np.asarray(reference_depth, dtype=np.float32), + **attributes, + ) + texture_path = output_directory / "texture.png" + with Image.open(image_path) as image: + image.convert("RGBA").save(texture_path, format="PNG", compress_level=1) + return result_path, texture_path + + def depth_edge_mask(depth, threshold=12.0 / 255.0): import numpy as np - padded = np.pad(depth.astype(np.float32), 1, mode="edge") - grad_x = ( - -padded[:-2, :-2] - + padded[:-2, 2:] - - 2.0 * padded[1:-1, :-2] - + 2.0 * padded[1:-1, 2:] - - padded[2:, :-2] - + padded[2:, 2:] - ) - grad_y = ( - -padded[:-2, :-2] - - 2.0 * padded[:-2, 1:-1] - - padded[:-2, 2:] - + padded[2:, :-2] - + 2.0 * padded[2:, 1:-1] - + padded[2:, 2:] - ) - magnitude = np.hypot(grad_x, grad_y) - finite = np.isfinite(magnitude) - if not finite.any(): - return np.ones(depth.shape, dtype=bool) - minimum = float(magnitude[finite].min()) - maximum = float(magnitude[finite].max()) - if maximum <= minimum: - return np.zeros(depth.shape, dtype=bool) - normalized = (magnitude - minimum) / (maximum - minimum) - return normalized >= threshold + values = np.asarray(depth, dtype=np.float32) + edges = np.zeros(values.shape, dtype=bool) + for axis in (0, 1): + first = np.take(values, range(values.shape[axis] - 1), axis=axis) + second = np.take(values, range(1, values.shape[axis]), axis=axis) + valid = ( + np.isfinite(first) + & np.isfinite(second) + & (first > 0.0) + & (second > 0.0) + ) + scale = np.maximum(np.minimum(np.abs(first), np.abs(second)), 1e-6) + discontinuity = valid & (np.abs(first - second) / scale > threshold) + lower = [slice(None), slice(None)] + upper = [slice(None), slice(None)] + lower[axis] = slice(0, -1) + upper[axis] = slice(1, None) + edges[tuple(lower)] |= discontinuity + edges[tuple(upper)] |= discontinuity + return edges + + +def edge_culling_stretch_limit(culling): + factor = min(max(float(culling) / 100.0, 0.0001), 1.0) + return 1.0 / factor + + +def foreground_mesh_mask(foreground_path, depth_shape): + import numpy as np + from PIL import Image, ImageFilter + + height, width = depth_shape + with Image.open(foreground_path) as source: + alpha = source.convert("RGBA").getchannel("A") + alpha = alpha.resize( + (width, height), + resample=Image.Resampling.BILINEAR, + ) + alpha = alpha.filter(ImageFilter.GaussianBlur(radius=0.75)) + return np.asarray(alpha) >= 128 def confidence_percentile_threshold(depth, confidence, percentile): import numpy as np + if float(percentile) <= 0.0: + return None valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence) if not valid.any(): raise ValueError("DA3 returned no valid depth confidence values") @@ -416,6 +739,9 @@ def mesh_triangle_faces( confidence_threshold, filter_edges, max_edge_length, + edge_threshold=12.0 / 255.0, + edge_stretch_limit=None, + vertex_mask=None, ): import numpy as np @@ -434,24 +760,47 @@ def mesh_triangle_faces( np.isfinite(points).all(axis=-1) & np.isfinite(depth) & (depth > 0.0) - & (confidence >= confidence_threshold) ) + if confidence_threshold is not None: + valid &= np.isfinite(confidence) & ( + confidence >= confidence_threshold + ) if filter_edges: - valid &= ~depth_edge_mask(depth) - face_valid = valid.reshape(-1)[quads].all(axis=1) - if max_edge_length > 0.0: - flat_points = points.reshape(-1, 3) - quad_points = flat_points[quads] - next_points = np.roll(quad_points, -1, axis=1) - edge_lengths = np.linalg.norm(quad_points - next_points, axis=-1) - face_valid &= (edge_lengths <= max_edge_length).all(axis=1) - quads = quads[face_valid] - if not len(quads): - return np.empty((0, 3), dtype=np.int32) - return np.concatenate( + valid &= ~depth_edge_mask(depth, edge_threshold) + if vertex_mask is not None: + if np.shape(vertex_mask) != depth.shape: + raise ValueError("Foreground mask and depth dimensions do not match") + valid &= np.asarray(vertex_mask, dtype=bool) + triangles = np.concatenate( (quads[:, (0, 1, 2)], quads[:, (0, 2, 3)]), axis=0, ) + face_valid = valid.reshape(-1)[triangles].all(axis=1) + if edge_stretch_limit is not None: + face_valid &= triangle_edge_stretch_mask( + points, + triangles, + edge_stretch_limit, + ) + if max_edge_length > 0.0: + flat_points = points.reshape(-1, 3) + triangle_points = flat_points[triangles] + next_points = np.roll(triangle_points, -1, axis=1) + edge_lengths = np.linalg.norm( + triangle_points - next_points, + axis=-1, + ) + face_valid &= (edge_lengths <= max_edge_length).all(axis=1) + triangles = triangles[face_valid] + if not len(triangles): + return np.empty((0, 3), dtype=np.int32) + return triangles + + +def triangle_edge_stretch_mask(points, triangles, stretch_limit): + stretch = triangle_edge_stretch_values(points, triangles) + tolerance = float(stretch_limit) * (1.0 + 1e-5) + return (stretch <= tolerance).all(axis=1) def export_textured_mesh_glb( @@ -512,6 +861,16 @@ def export_textured_mesh_glb( visual=visual, process=False, ) + mesh.vertex_attributes["o_depth"] = depth[index].reshape(-1) + mesh.vertex_attributes["o_confidence"] = confidence[ + index + ].reshape(-1) + mesh.vertex_attributes["o_confidence_percentile"] = ( + confidence_percentile_map( + depth[index], + confidence[index], + ).reshape(-1) + ) mesh.remove_unreferenced_vertices() meshes.append(mesh) referenced_points.append(mesh.vertices) @@ -1160,6 +1519,9 @@ def exported_paths(output_directory, export_format): if export_format == "mesh": result_path = output_directory / "mesh.glb" return [result_path] if result_path.is_file() else [] + if export_format == "mesh_data": + result_path = output_directory / "mesh.npz" + return [result_path] if result_path.is_file() else [] if export_format == "mini_npz": result_path = output_directory / "results.npz" if result_path.is_file(): @@ -1401,6 +1763,23 @@ def reconstruct(args, model_cache=None, cancel_check=None): bool(getattr(args, "mesh_filter_edges", True)), float(getattr(args, "mesh_max_edge_length", 0.1)), ) + elif args.export_format == "mesh_data": + def mesh_progress(step, total, label): + write_status( + args.status, + 0.88 + 0.09 * step / total, + f"Post-processing {step}/{total}: {label}", + ) + + export_image_mesh_data( + prediction, + images[0], + output_directory, + foreground_path=( + foreground_images[0] if foreground_images else None + ), + progress_callback=mesh_progress, + ) else: export( prediction, @@ -1408,6 +1787,13 @@ def reconstruct(args, model_cache=None, cancel_check=None): str(output_directory), **export_kwargs, ) + if args.export_format in {"glb", "gs_ply", "depth_vis", "mesh"}: + write_status( + args.status, + 0.97, + "Writing structured debug attributes", + ) + export_prediction_attributes(prediction, output_directory) check_cancelled(cancel_check) result_paths = wait_for_exported_paths( diff --git a/tests/test_addon.py b/tests/test_addon.py index ac17a55..b740689 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -8,6 +8,8 @@ from tempfile import TemporaryDirectory from types import ModuleType, SimpleNamespace from unittest.mock import patch +import numpy as np + PROJECT_ROOT = Path(__file__).resolve().parents[1] SOURCE_ROOT = PROJECT_ROOT / "src" @@ -28,6 +30,7 @@ class AddonTest(unittest.TestCase): ) fake_bpy.types = SimpleNamespace( AddonPreferences=object, + Menu=object, Operator=object, Panel=object, PropertyGroup=object, @@ -79,6 +82,7 @@ class AddonTest(unittest.TestCase): self.context_menu_draws.clear() self.addon.operators._active_job = None self.addon.image_depth._last_depth_plane_model = "BASE" + self.addon.image_mesh._last_mesh_model = "BASE" self.fake_bpy.data = SimpleNamespace( objects=SimpleNamespace(get=lambda _name: None), ) @@ -125,6 +129,135 @@ class AddonTest(unittest.TestCase): ) self.assertFalse(hasattr(self.addon, "ImportDepthPlane")) + def test_registers_perspective_mesh_operators(self): + self.assertIn(self.addon.ConvertImageToMesh, self.addon.CLASSES) + self.assertIn(self.addon.BuildImageMesh, self.addon.CLASSES) + self.assertEqual(self.addon.BuildImageMesh.bl_options, {"INTERNAL"}) + self.assertIn(self.addon.SnapFormImageMenu, self.addon.CLASSES) + + def test_image_context_menu_groups_actions_and_separates_background(self): + calls = [] + + class Layout: + operator_context = None + + def operator(self, identifier, **_options): + calls.append(identifier) + + def separator(self): + calls.append("separator") + + menu = self.addon.SnapFormImageMenu() + menu.layout = Layout() + menu.draw(None) + + self.assertEqual( + calls, + [ + "snapform.convert_image_to_depth_plane", + "snapform.convert_image_to_mesh", + "separator", + "snapform.remove_image_background", + ], + ) + + def test_image_mesh_reference_plane_matches_image_empty_bounds(self): + source_object = SimpleNamespace( + data=SimpleNamespace(size=(401, 201)), + empty_display_size=4.0, + empty_image_offset=(-0.25, -0.75), + ) + intrinsics = np.asarray( + [[200.0, 0.0, 200.0], [0.0, 200.0, 100.0], [0.0, 0.0, 1.0]], + dtype=np.float32, + ) + camera_vertices = np.asarray( + [ + [-2.0, -1.0, 2.0], + [2.0, -1.0, 2.0], + [-2.0, 1.0, 2.0], + [2.0, 1.0, 2.0], + ], + dtype=np.float32, + ) + + vertices, camera = self.addon.image_mesh.empty_anchored_vertices( + source_object, + camera_vertices, + intrinsics, + (401, 201), + 2.0, + ) + + x_min, x_max, y_min, y_max = self.addon.image_depth.plane_bounds( + source_object + ) + np.testing.assert_allclose( + vertices, + [ + [x_min, y_max, 0.0], + [x_max, y_max, 0.0], + [x_min, y_min, 0.0], + [x_max, y_min, 0.0], + ], + ) + np.testing.assert_allclose( + camera["location"], + ((x_min + x_max) * 0.5, (y_min + y_max) * 0.5, 2.0024922), + ) + + def test_mesh_conversion_dialog_only_keeps_model_and_camera_options(self): + annotations = self.addon.ConvertImageToMesh.__annotations__ + + for property_name in ( + "snapform_confidence_percentile", + "snapform_edge_culling", + "snapform_face_reduction", + "snapform_depth_smoothing", + "snapform_remove_background", + "snapform_postprocess_method", + ): + with self.subTest(property_name=property_name): + self.assertNotIn(property_name, annotations) + self.assertFalse(annotations["snapform_create_camera"]["default"]) + self.assertFalse(annotations["snapform_unshaded"]["default"]) + self.assertEqual(annotations["snapform_unshaded"]["name"], "Unshaded") + + def test_image_mesh_uses_o_prefixed_attributes_and_bundled_nodes(self): + image_mesh = self.addon.image_mesh + names = set(image_mesh.POINT_ATTRIBUTES) | set( + image_mesh.FACE_ATTRIBUTES + ) + + self.assertTrue(all(name.startswith("o_") for name in names)) + self.assertTrue(image_mesh.node_asset_path().is_file()) + self.assertEqual(image_mesh.NODE_GROUP_NAME, "O Refine DA3 Mesh") + self.assertEqual(image_mesh.NODE_ASSET_NAME, "O_SnapForm.blend") + + def test_image_mesh_material_depth_scale_matches_fitted_geometry(self): + mesh_object = SimpleNamespace(scale=SimpleNamespace(z=-2.0)) + + self.assertEqual( + self.addon.image_mesh.material_depth_scale( + mesh_object, + {"scale_z": 0.25}, + ), + 0.5, + ) + + def test_image_mesh_default_resolution_uses_double_reduction_estimate(self): + vertices = np.asarray( + ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), + dtype=np.float32, + ) + faces = np.asarray(((0, 1, 2),), dtype=np.int32) + edge_lengths = np.asarray((1.0, 1.0, np.sqrt(2.0))) + + self.assertAlmostEqual( + self.addon.image_mesh.default_remesh_resolution(vertices, faces), + float(np.median(edge_lengths) * np.sqrt(10.0) * 2.0), + ) + def test_registers_image_background_operators(self): self.assertIn(self.addon.RemoveImageBackground, self.addon.CLASSES) self.assertIn(self.addon.ApplyRemovedBackground, self.addon.CLASSES) @@ -1486,6 +1619,51 @@ class AddonTest(unittest.TestCase): ) self.assertEqual(parameters["use_ray_pose"], 0) + def test_reconstruct_request_enables_ben2_for_image_mesh(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()), + ): + request = runtime.reconstruct_request( + settings, + "mesh_data", + Path("status.json"), + Path("result"), + remove_background=True, + ) + + parameters = request.payload["parameters"] + self.assertNotIn("mesh_confidence_percentile", parameters) + self.assertNotIn("mesh_edge_culling", parameters) + self.assertNotIn("mesh_face_reduction", parameters) + self.assertNotIn("mesh_depth_smoothing", parameters) + self.assertNotIn("mesh_postprocess_method", parameters) + self.assertEqual( + parameters["remove_background_model_dir"], + str(Path("ben2_base")), + ) + def test_export_format_property_includes_all_da3_formats(self): export_format = self.addon.SnapFormSettings.__annotations__[ "export_format" diff --git a/tests/test_image_depth.py b/tests/test_image_depth.py index 32c21d3..f9d14aa 100644 --- a/tests/test_image_depth.py +++ b/tests/test_image_depth.py @@ -21,7 +21,7 @@ def load_image_depth(): IntProperty=lambda **options: options, StringProperty=lambda **options: options, ) - fake_bpy.types = SimpleNamespace(Operator=object) + fake_bpy.types = SimpleNamespace(Menu=object, Operator=object) fake_bpy.path = SimpleNamespace(abspath=lambda path: path) with patch.dict(sys.modules, {"bpy": fake_bpy}): specification = importlib.util.spec_from_file_location( diff --git a/tests/test_pack.py b/tests/test_pack.py index 8b0c14c..e2f0a72 100644 --- a/tests/test_pack.py +++ b/tests/test_pack.py @@ -56,6 +56,10 @@ class PackTest(unittest.TestCase): self.assertIn("__init__.py", names) self.assertIn("blender_manifest.toml", names) self.assertIn("image_depth.py", names) + self.assertIn( + "assets/O_SnapForm.blend", + names, + ) self.assertIn("worker.py", names) self.assertIn("ai_host.py", names) self.assertNotIn("snapform/__init__.py", names) diff --git a/tests/test_worker.py b/tests/test_worker.py index 05ab409..a7655fe 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -982,7 +982,14 @@ class WorkerTest(unittest.TestCase): "export_textured_mesh_glb", side_effect=fake_mesh_export, ): - for export_format in self.worker.EXPORT_FORMATS: + debug_formats = ( + "glb", + "gs_ply", + "depth_vis", + "mini_npz", + "mesh", + ) + for export_format in debug_formats: args = SimpleNamespace( input=str(input_path), video_frames=16, @@ -1172,6 +1179,124 @@ class WorkerTest(unittest.TestCase): self.assertEqual(threshold, 2.5) + def test_zero_mesh_confidence_disables_confidence_filter(self): + depth = np.ones((2, 2), dtype=np.float32) + confidence = np.full((2, 2), np.nan, dtype=np.float32) + x, y = np.meshgrid([0.0, 1.0], [0.0, 1.0]) + points = np.stack((x, y, depth), axis=-1) + + threshold = self.worker.confidence_percentile_threshold( + depth, + confidence, + 0.0, + ) + faces = self.worker.mesh_triangle_faces( + points, + depth, + confidence, + threshold, + filter_edges=False, + max_edge_length=0.0, + ) + + self.assertIsNone(threshold) + self.assertEqual(len(faces), 2) + + def test_mesh_filters_triangles_independently(self): + depth = np.ones((2, 2), dtype=np.float32) + confidence = np.ones((2, 2), dtype=np.float32) + confidence[1, 0] = 0.0 + x, y = np.meshgrid( + np.arange(2, dtype=np.float32), + np.arange(2, dtype=np.float32), + ) + points = np.stack((x, y, depth), axis=-1) + + faces = self.worker.mesh_triangle_faces( + points, + depth, + confidence, + confidence_threshold=0.5, + filter_edges=False, + max_edge_length=0.0, + ) + + np.testing.assert_array_equal(faces, [[0, 3, 1]]) + + def test_mesh_background_mask_combines_with_confidence(self): + depth = np.ones((2, 3), dtype=np.float32) + confidence = np.ones((2, 3), dtype=np.float32) + x, y = np.meshgrid( + np.arange(3, dtype=np.float32), + np.arange(2, dtype=np.float32), + ) + points = np.stack((x, y, depth), axis=-1) + foreground = np.ones((2, 3), dtype=bool) + foreground[:, 0] = False + + faces = self.worker.mesh_triangle_faces( + points, + depth, + confidence, + confidence_threshold=None, + filter_edges=False, + max_edge_length=0.0, + vertex_mask=foreground, + ) + + np.testing.assert_array_equal(faces, [[1, 4, 5], [1, 5, 2]]) + + def test_depth_edge_filter_ignores_smooth_slope(self): + smooth = np.tile( + np.linspace(2.0, 2.04, 5, dtype=np.float32), + (3, 1), + ) + discontinuous = smooth.copy() + discontinuous[:, 3:] += 1.0 + + self.assertFalse(self.worker.depth_edge_mask(smooth).any()) + self.assertTrue(self.worker.depth_edge_mask(discontinuous)[:, 2:4].all()) + + def test_higher_edge_culling_allows_less_stretch(self): + limits = [ + self.worker.edge_culling_stretch_limit(value) + for value in (0.0, 50.0, 100.0) + ] + + self.assertEqual(limits, [10000.0, 2.0, 1.0]) + + def test_stretched_edge_culling_uses_camera_space_edge_ratio(self): + flat_points = np.asarray( + [ + [[0.0, 0.0, 1.0], [1.0, 0.0, 1.0]], + [[0.0, 1.0, 1.0], [1.0, 1.0, 1.0]], + ], + dtype=np.float32, + ) + stretched_points = flat_points.copy() + stretched_points[0, 1] = (2.0, 0.0, 2.0) + triangles = np.asarray([[0, 2, 3], [0, 3, 1]], dtype=np.int32) + + flat_mask = self.worker.triangle_edge_stretch_mask( + flat_points, + triangles, + 1.0, + ) + loose_mask = self.worker.triangle_edge_stretch_mask( + stretched_points, + triangles, + 2.0, + ) + strict_mask = self.worker.triangle_edge_stretch_mask( + stretched_points, + triangles, + 1.0, + ) + + self.assertTrue(flat_mask.all()) + self.assertTrue(loose_mask.all()) + self.assertFalse(strict_mask.all()) + def test_mesh_faces_apply_confidence_threshold_and_face_direction(self): points = np.array( [ @@ -1228,6 +1353,113 @@ class WorkerTest(unittest.TestCase): self.assertEqual(face_counts, [8, 2, 0]) + def test_mesh_compaction_removes_unreferenced_vertices(self): + vertices = np.arange(18, dtype=np.float32).reshape(6, 3) + faces = np.asarray([[1, 3, 5]], dtype=np.int32) + + compact_vertices, compact_faces = self.worker.compact_mesh( + vertices, + faces, + ) + + np.testing.assert_array_equal( + compact_vertices, + vertices[[1, 3, 5]], + ) + np.testing.assert_array_equal(compact_faces, [[0, 1, 2]]) + + def test_image_mesh_data_keeps_camera_geometry_and_uv(self): + 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): + Path(path).write_bytes(b"texture") + + fake_pil = ModuleType("PIL") + fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage()) + depth = np.full((1, 3, 3), 2.0, dtype=np.float32) + confidence = np.ones((1, 3, 3), dtype=np.float32) + intrinsics = np.asarray( + [[[2.0, 0.0, 1.0], [0.0, 2.0, 1.0], [0.0, 0.0, 1.0]]], + dtype=np.float32, + ) + prediction = SimpleNamespace( + depth=depth, + conf=confidence, + intrinsics=intrinsics, + ) + progress = [] + + with tempfile.TemporaryDirectory() as directory, patch.dict( + sys.modules, + {"PIL": fake_pil}, + ), patch.object( + self.worker, + "compact_mesh", + side_effect=lambda vertices, faces: ( + vertices, + faces, + ), + ): + result_path, texture_path = self.worker.export_image_mesh_data( + prediction, + Path(directory) / "input.png", + Path(directory) / "result", + progress_callback=lambda step, total, label: progress.append( + (step, total, label) + ), + ) + with np.load(result_path) as result: + self.assertEqual(result["vertices"].shape, (9, 3)) + self.assertEqual(result["faces"].shape, (8, 3)) + self.assertEqual(result["uv"].shape, (9, 2)) + self.assertEqual(float(result["reference_depth"]), 2.0) + self.assertNotIn("postprocess_method", result.files) + self.assertEqual(result["point_depth"].shape, (9,)) + self.assertEqual( + result["point_confidence_percentile"].shape, + (9,), + ) + self.assertEqual( + result["face_edge_stretch"].shape, + (8,), + ) + self.assertEqual(result["face_background"].shape, (8,)) + np.testing.assert_allclose(result["uv"][0], [0.0, 1.0]) + np.testing.assert_allclose(result["uv"][-1], [1.0, 0.0]) + self.assertTrue(texture_path.is_file()) + self.assertEqual([item[0] for item in progress], list(range(1, 7))) + self.assertTrue(all(item[1] == 6 for item in progress)) + + def test_prediction_attributes_keep_debug_structure(self): + prediction = SimpleNamespace( + depth=np.ones((1, 2, 3), dtype=np.float32), + conf=np.arange(6, dtype=np.float32).reshape(1, 2, 3), + intrinsics=np.eye(3, dtype=np.float32)[None], + extrinsics=np.eye(4, dtype=np.float32)[None], + ) + + with tempfile.TemporaryDirectory() as directory: + output_path = self.worker.export_prediction_attributes( + prediction, + directory, + ) + with np.load(output_path) as attributes: + self.assertEqual(attributes["depth"].shape, (1, 2, 3)) + self.assertEqual(attributes["conf"].shape, (1, 2, 3)) + self.assertEqual( + attributes["confidence_percentile"].shape, + (1, 2, 3), + ) + if __name__ == "__main__": unittest.main() diff --git a/tools/build_geometry_nodes_asset.py b/tools/build_geometry_nodes_asset.py new file mode 100644 index 0000000..03d3053 --- /dev/null +++ b/tools/build_geometry_nodes_asset.py @@ -0,0 +1,633 @@ +from pathlib import Path + +import bpy + + +GROUP_NAME = "O Refine DA3 Mesh" +OUTPUT_PATH = ( + Path(__file__).parents[1] + / "src" + / "snapform" + / "assets" + / "O_SnapForm.blend" +) + + +def input_socket(node, name, fallback): + value = node.inputs.get(name) + return value if value is not None else node.inputs[fallback] + + +def new_socket(interface, name, in_out, socket_type, parent=None): + return interface.new_socket( + name=name, + in_out=in_out, + socket_type=socket_type, + parent=parent, + ) + + +def new_compare(nodes, data_type, operation, location, value): + node = nodes.new("FunctionNodeCompare") + node.data_type = data_type + node.operation = operation + node.location = location + input_socket(node, "B", 1).default_value = value + return node + + +def new_group_input(nodes, location, visible_names): + node = nodes.new("NodeGroupInput") + node.location = location + visible = set(visible_names) + for output in node.outputs: + output.hide = output.name not in visible + return node + + +def add_edge_boundary_field(nodes, links, location): + x, y = location + vertex_neighbors = nodes.new("GeometryNodeInputMeshVertexNeighbors") + vertex_neighbors.location = (x, y + 100.0) + loose_vertex = new_compare( + nodes, + "INT", + "EQUAL", + (x + 180.0, y + 100.0), + 1, + ) + links.new( + vertex_neighbors.outputs["Vertex Count"], + input_socket(loose_vertex, "A", 0), + ) + edge_neighbors = nodes.new("GeometryNodeInputMeshEdgeNeighbors") + edge_neighbors.location = (x, y - 100.0) + boundary_edge = new_compare( + nodes, + "INT", + "EQUAL", + (x + 180.0, y - 100.0), + 1, + ) + links.new( + edge_neighbors.outputs["Face Count"], + input_socket(boundary_edge, "A", 0), + ) + edge_to_point = nodes.new("GeometryNodeFieldOnDomain") + edge_to_point.data_type = "BOOLEAN" + edge_to_point.domain = "EDGE" + edge_to_point.location = (x + 370.0, y - 100.0) + links.new(boundary_edge.outputs["Result"], edge_to_point.inputs["Value"]) + combine = nodes.new("FunctionNodeBooleanMath") + combine.operation = "OR" + combine.location = (x + 560.0, y) + links.new(loose_vertex.outputs["Result"], combine.inputs[0]) + links.new(edge_to_point.outputs["Value"], combine.inputs[1]) + return combine.outputs[0] + + +def add_topology_fix_zone(group, geometry, iterations, location): + nodes = group.nodes + links = group.links + x, y = location + repeat_output = nodes.new("GeometryNodeRepeatOutput") + repeat_output.location = (x + 1450.0, y) + repeat_input = nodes.new("GeometryNodeRepeatInput") + repeat_input.pair_with_output(repeat_output) + repeat_input.location = (x, y) + links.new(geometry, repeat_input.inputs["Geometry"]) + links.new(iterations, repeat_input.inputs["Iterations"]) + + edge_neighbors = nodes.new("GeometryNodeInputMeshEdgeNeighbors") + edge_neighbors.location = (x + 100.0, y + 420.0) + non_manifold = new_compare( + nodes, + "INT", + "GREATER_THAN", + (x + 280.0, y + 500.0), + 2, + ) + boundary = new_compare( + nodes, + "INT", + "LESS_THAN", + (x + 280.0, y + 320.0), + 2, + ) + links.new( + edge_neighbors.outputs["Face Count"], + input_socket(non_manifold, "A", 0), + ) + links.new( + edge_neighbors.outputs["Face Count"], + input_socket(boundary, "A", 0), + ) + non_manifold_face = nodes.new("GeometryNodeFieldOnDomain") + non_manifold_face.data_type = "BOOLEAN" + non_manifold_face.domain = "FACE" + non_manifold_face.location = (x + 480.0, y + 500.0) + boundary_face = nodes.new("GeometryNodeFieldOnDomain") + boundary_face.data_type = "BOOLEAN" + boundary_face.domain = "FACE" + boundary_face.location = (x + 480.0, y + 320.0) + links.new(non_manifold.outputs["Result"], non_manifold_face.inputs["Value"]) + links.new(boundary.outputs["Result"], boundary_face.inputs["Value"]) + broken_face = nodes.new("FunctionNodeBooleanMath") + broken_face.operation = "AND" + broken_face.location = (x + 670.0, y + 420.0) + links.new(non_manifold_face.outputs["Value"], broken_face.inputs[0]) + links.new(boundary_face.outputs["Value"], broken_face.inputs[1]) + remove_broken = nodes.new("GeometryNodeSeparateGeometry") + remove_broken.domain = "FACE" + remove_broken.location = (x + 670.0, y) + links.new(repeat_input.outputs["Geometry"], remove_broken.inputs["Geometry"]) + links.new(broken_face.outputs[0], remove_broken.inputs["Selection"]) + + remaining_neighbors = nodes.new("GeometryNodeInputMeshEdgeNeighbors") + remaining_neighbors.location = (x + 650.0, y - 420.0) + remaining_boundary = new_compare( + nodes, + "INT", + "LESS_THAN", + (x + 830.0, y - 420.0), + 2, + ) + links.new( + remaining_neighbors.outputs["Face Count"], + input_socket(remaining_boundary, "A", 0), + ) + boundary_face_2 = nodes.new("GeometryNodeFieldOnDomain") + boundary_face_2.data_type = "BOOLEAN" + boundary_face_2.domain = "FACE" + boundary_face_2.location = (x + 1010.0, y - 420.0) + boundary_edge_2 = nodes.new("GeometryNodeFieldOnDomain") + boundary_edge_2.data_type = "BOOLEAN" + boundary_edge_2.domain = "EDGE" + boundary_edge_2.location = (x + 1190.0, y - 420.0) + links.new(remaining_boundary.outputs["Result"], boundary_face_2.inputs["Value"]) + links.new(boundary_face_2.outputs["Value"], boundary_edge_2.inputs["Value"]) + face_groups = nodes.new("GeometryNodeEdgesToFaceGroups") + face_groups.location = (x + 1010.0, y - 220.0) + links.new( + boundary_edge_2.outputs["Value"], + face_groups.inputs["Boundary Edges"], + ) + accumulate = nodes.new("GeometryNodeAccumulateField") + accumulate.data_type = "INT" + accumulate.domain = "FACE" + accumulate.location = (x + 1190.0, y - 160.0) + links.new(remaining_neighbors.outputs["Face Count"], accumulate.inputs["Value"]) + links.new(face_groups.outputs["Face Group ID"], accumulate.inputs["Group ID"]) + single_face = new_compare( + nodes, + "INT", + "LESS_THAN", + (x + 1370.0, y - 160.0), + 2, + ) + links.new(accumulate.outputs["Total"], input_socket(single_face, "A", 0)) + remove_single = nodes.new("GeometryNodeSeparateGeometry") + remove_single.domain = "FACE" + remove_single.location = (x + 1250.0, y + 80.0) + links.new(remove_broken.outputs["Inverted"], remove_single.inputs["Geometry"]) + links.new(single_face.outputs["Result"], remove_single.inputs["Selection"]) + links.new(remove_single.outputs["Inverted"], repeat_output.inputs["Geometry"]) + return repeat_output.outputs["Geometry"] + + +def add_smooth_step(group, geometry, amount, location): + nodes = group.nodes + links = group.links + x, y = location + position = nodes.new("GeometryNodeInputPosition") + position.location = (x + 100.0, y + 500.0) + smooth_all = nodes.new("GeometryNodeBlurAttribute") + smooth_all.data_type = "FLOAT_VECTOR" + smooth_all.location = (x + 300.0, y + 500.0) + smooth_all.inputs["Iterations"].default_value = 1 + smooth_all.inputs["Weight"].default_value = 1.0 + links.new(position.outputs["Position"], smooth_all.inputs["Value"]) + + boundary_field = add_edge_boundary_field( + nodes, + links, + (x + 100.0, y - 350.0), + ) + boundary_on_points = nodes.new("GeometryNodeFieldOnDomain") + boundary_on_points.data_type = "BOOLEAN" + boundary_on_points.domain = "POINT" + boundary_on_points.location = (x + 700.0, y - 350.0) + links.new(boundary_field, boundary_on_points.inputs["Value"]) + boundary_geometry = nodes.new("GeometryNodeSeparateGeometry") + boundary_geometry.domain = "POINT" + boundary_geometry.location = (x + 880.0, y - 220.0) + links.new(geometry, boundary_geometry.inputs["Geometry"]) + links.new(boundary_on_points.outputs["Value"], boundary_geometry.inputs["Selection"]) + + smooth_boundary = nodes.new("GeometryNodeBlurAttribute") + smooth_boundary.data_type = "FLOAT_VECTOR" + smooth_boundary.location = (x + 500.0, y + 300.0) + smooth_boundary.inputs["Iterations"].default_value = 1 + smooth_boundary.inputs["Weight"].default_value = 0.5 + links.new(position.outputs["Position"], smooth_boundary.inputs["Value"]) + nearest_boundary = nodes.new("GeometryNodeSampleNearest") + nearest_boundary.domain = "POINT" + nearest_boundary.location = (x + 1050.0, y - 120.0) + links.new(boundary_geometry.outputs["Selection"], nearest_boundary.inputs["Geometry"]) + sample_boundary = nodes.new("GeometryNodeSampleIndex") + sample_boundary.data_type = "FLOAT_VECTOR" + sample_boundary.domain = "POINT" + sample_boundary.clamp = False + sample_boundary.location = (x + 1220.0, y + 20.0) + links.new(boundary_geometry.outputs["Selection"], sample_boundary.inputs["Geometry"]) + links.new(nearest_boundary.outputs["Index"], sample_boundary.inputs["Index"]) + links.new(smooth_boundary.outputs["Value"], sample_boundary.inputs["Value"]) + + boundary_switch = nodes.new("GeometryNodeSwitch") + boundary_switch.input_type = "VECTOR" + boundary_switch.location = (x + 1050.0, y + 350.0) + links.new(boundary_on_points.outputs["Value"], boundary_switch.inputs["Switch"]) + links.new(smooth_all.outputs["Value"], boundary_switch.inputs["False"]) + links.new(sample_boundary.outputs["Value"], boundary_switch.inputs["True"]) + + normal = nodes.new("GeometryNodeInputNormal") + normal.location = (x + 300.0, y + 100.0) + blur_normal = nodes.new("GeometryNodeBlurAttribute") + blur_normal.data_type = "FLOAT_VECTOR" + blur_normal.location = (x + 500.0, y + 100.0) + blur_normal.inputs["Iterations"].default_value = 10 + links.new(normal.outputs["Normal"], blur_normal.inputs["Value"]) + normal_length = nodes.new("ShaderNodeVectorMath") + normal_length.operation = "LENGTH" + normal_length.location = (x + 700.0, y + 100.0) + links.new(blur_normal.outputs["Value"], normal_length.inputs[0]) + weight = nodes.new("ShaderNodeMath") + weight.operation = "MULTIPLY" + weight.location = (x + 1050.0, y + 150.0) + links.new(normal_length.outputs["Value"], weight.inputs[0]) + links.new(amount, weight.inputs[1]) + + mix = nodes.new("ShaderNodeMix") + mix.data_type = "VECTOR" + mix.location = (x + 1270.0, y + 300.0) + links.new(weight.outputs[0], input_socket(mix, "Factor", 0)) + links.new(position.outputs["Position"], input_socket(mix, "A", 4)) + links.new(boundary_switch.outputs["Output"], input_socket(mix, "B", 5)) + set_position = nodes.new("GeometryNodeSetPosition") + set_position.location = (x + 1450.0, y) + links.new(geometry, set_position.inputs["Geometry"]) + links.new(mix.outputs["Result"], set_position.inputs["Position"]) + return set_position.outputs["Geometry"] + + +def add_remesh_zone( + group, + geometry, + resolution, + iterations, + fix_steps, + smooth_amount, + location, +): + nodes = group.nodes + links = group.links + x, y = location + repeat_output = nodes.new("GeometryNodeRepeatOutput") + repeat_output.location = (x + 3300.0, y) + repeat_input = nodes.new("GeometryNodeRepeatInput") + repeat_input.pair_with_output(repeat_output) + repeat_input.location = (x, y) + links.new(geometry, repeat_input.inputs["Geometry"]) + links.new(iterations, repeat_input.inputs["Iterations"]) + + half_resolution = nodes.new("ShaderNodeMath") + half_resolution.operation = "MULTIPLY" + half_resolution.inputs[1].default_value = 0.5 + half_resolution.location = (x + 100.0, y - 250.0) + links.new(resolution, half_resolution.inputs[0]) + merge = nodes.new("GeometryNodeMergeByDistance") + merge.mode = "CONNECTED" + merge.location = (x + 300.0, y) + links.new(repeat_input.outputs["Geometry"], merge.inputs["Geometry"]) + links.new(half_resolution.outputs[0], merge.inputs["Distance"]) + fixed = add_topology_fix_zone( + group, + merge.outputs["Geometry"], + fix_steps, + (x + 550.0, y), + ) + smoothed = add_smooth_step( + group, + fixed, + smooth_amount, + (x + 2100.0, y), + ) + links.new(smoothed, repeat_output.inputs["Geometry"]) + return repeat_output.outputs["Geometry"] + + +def add_uv_reprojection(group, source, geometry, location): + nodes = group.nodes + links = group.links + x, y = location + position = nodes.new("GeometryNodeInputPosition") + position.location = (x, y - 200.0) + source_uv = nodes.new("GeometryNodeInputNamedAttribute") + source_uv.data_type = "FLOAT_VECTOR" + source_uv.inputs["Name"].default_value = "UVMap" + source_uv.location = (x, y) + sample_uv = nodes.new("GeometryNodeSampleNearestSurface") + sample_uv.data_type = "FLOAT_VECTOR" + sample_uv.location = (x + 430.0, y) + links.new(source, sample_uv.inputs["Mesh"]) + links.new(source_uv.outputs["Attribute"], sample_uv.inputs["Value"]) + links.new(position.outputs["Position"], sample_uv.inputs["Sample Position"]) + store_uv = nodes.new("GeometryNodeStoreNamedAttribute") + store_uv.data_type = "FLOAT2" + store_uv.domain = "CORNER" + store_uv.location = (x + 1050.0, y + 250.0) + store_uv.inputs["Name"].default_value = "UVMap" + links.new(geometry, store_uv.inputs["Geometry"]) + links.new(sample_uv.outputs["Value"], store_uv.inputs["Value"]) + return store_uv.outputs["Geometry"] + + +def build_group(): + group = bpy.data.node_groups.new(GROUP_NAME, "GeometryNodeTree") + group.is_modifier = True + interface = group.interface + new_socket(interface, "Geometry", "INPUT", "NodeSocketGeometry") + + resolution = new_socket( + interface, + "Resolution", + "INPUT", + "NodeSocketFloat", + ) + resolution.default_value = 0.01 + resolution.min_value = 0.000001 + resolution.subtype = "DISTANCE" + remesh_iterations = new_socket( + interface, + "Relax Iterations", + "INPUT", + "NodeSocketInt", + ) + remesh_iterations.default_value = 2 + remesh_iterations.min_value = 1 + remesh_iterations.max_value = 8 + smooth_amount = new_socket( + interface, + "Smooth", + "INPUT", + "NodeSocketFloat", + ) + smooth_amount.default_value = 1.0 + smooth_amount.min_value = 0.0 + smooth_amount.max_value = 1.0 + smooth_amount.subtype = "FACTOR" + + culling_panel = interface.new_panel( + name="Culling", + default_closed=False, + ) + confidence = new_socket( + interface, + "Confidence", + "INPUT", + "NodeSocketFloat", + parent=culling_panel, + ) + confidence.default_value = 0.0 + confidence.min_value = 0.0 + confidence.max_value = 1.0 + confidence.subtype = "PERCENTAGE" + edge = new_socket( + interface, + "Stretch", + "INPUT", + "NodeSocketFloat", + parent=culling_panel, + ) + edge.default_value = 0.1 + edge.min_value = 0.0 + edge.max_value = 1.0 + edge.subtype = "PERCENTAGE" + background = new_socket( + interface, + "Background", + "INPUT", + "NodeSocketBool", + parent=culling_panel, + ) + background.default_value = True + + options_panel = interface.new_panel( + name="Options", + default_closed=True, + ) + fix_steps = new_socket( + interface, + "Fix Topology", + "INPUT", + "NodeSocketInt", + parent=options_panel, + ) + fix_steps.default_value = 3 + fix_steps.min_value = 1 + fix_steps.max_value = 8 + new_socket(interface, "Geometry", "OUTPUT", "NodeSocketGeometry") + + nodes = group.nodes + links = group.links + geometry_input = new_group_input(nodes, (-650.0, 50.0), {"Geometry"}) + confidence_input = new_group_input(nodes, (-1800.0, 800.0), {"Confidence"}) + stretch_input = new_group_input(nodes, (-1800.0, 50.0), {"Stretch"}) + background_input = new_group_input(nodes, (-1800.0, -350.0), {"Background"}) + remesh_input = new_group_input( + nodes, + (-100.0, -500.0), + {"Resolution", "Relax Iterations", "Smooth"}, + ) + topology_input = new_group_input(nodes, (350.0, -700.0), {"Fix Topology"}) + island_input = new_group_input(nodes, (3400.0, 100.0), {"Resolution"}) + uv_source_input = new_group_input(nodes, (4200.0, -300.0), {"Geometry"}) + group_output = nodes.new("NodeGroupOutput") + group_output.location = (5650.0, 100.0) + + confidence_attribute = nodes.new("GeometryNodeInputNamedAttribute") + confidence_attribute.data_type = "FLOAT" + confidence_attribute.inputs["Name"].default_value = "o_face_confidence" + confidence_attribute.location = (-1800.0, 600.0) + confidence_compare = new_compare( + nodes, + "FLOAT", + "LESS_THAN", + (-1550.0, 600.0), + 0.0, + ) + links.new( + confidence_attribute.outputs["Attribute"], + input_socket(confidence_compare, "A", 0), + ) + links.new( + confidence_input.outputs["Confidence"], + input_socket(confidence_compare, "B", 1), + ) + + edge_attribute = nodes.new("GeometryNodeInputNamedAttribute") + edge_attribute.data_type = "FLOAT" + edge_attribute.inputs["Name"].default_value = "o_edge_stretch" + edge_attribute.location = (-1800.0, 350.0) + edge_maximum = nodes.new("ShaderNodeMath") + edge_maximum.operation = "MAXIMUM" + edge_maximum.inputs[1].default_value = 0.0001 + edge_maximum.location = (-1600.0, 100.0) + edge_reciprocal = nodes.new("ShaderNodeMath") + edge_reciprocal.operation = "DIVIDE" + edge_reciprocal.inputs[0].default_value = 1.0 + edge_reciprocal.location = (-1400.0, 100.0) + edge_compare = new_compare( + nodes, + "FLOAT", + "GREATER_EQUAL", + (-1200.0, 350.0), + 1.0, + ) + links.new(stretch_input.outputs["Stretch"], edge_maximum.inputs[0]) + links.new(edge_maximum.outputs[0], edge_reciprocal.inputs[1]) + links.new(edge_attribute.outputs["Attribute"], input_socket(edge_compare, "A", 0)) + links.new(edge_reciprocal.outputs[0], input_socket(edge_compare, "B", 1)) + + background_attribute = nodes.new("GeometryNodeInputNamedAttribute") + background_attribute.data_type = "BOOLEAN" + background_attribute.inputs["Name"].default_value = "o_background_face" + background_attribute.location = (-1800.0, -200.0) + background_and = nodes.new("FunctionNodeBooleanMath") + background_and.operation = "AND" + background_and.location = (-1450.0, -200.0) + links.new(background_attribute.outputs["Attribute"], background_and.inputs[0]) + links.new(background_input.outputs["Background"], background_and.inputs[1]) + first_or = nodes.new("FunctionNodeBooleanMath") + first_or.operation = "OR" + first_or.location = (-950.0, 450.0) + links.new(confidence_compare.outputs["Result"], first_or.inputs[0]) + links.new(edge_compare.outputs["Result"], first_or.inputs[1]) + final_or = nodes.new("FunctionNodeBooleanMath") + final_or.operation = "OR" + final_or.location = (-750.0, 300.0) + links.new(first_or.outputs[0], final_or.inputs[0]) + links.new(background_and.outputs[0], final_or.inputs[1]) + full_stretch = new_compare( + nodes, + "FLOAT", + "GREATER_EQUAL", + (-950.0, -20.0), + 0.999999, + ) + links.new( + stretch_input.outputs["Stretch"], + input_socket(full_stretch, "A", 0), + ) + culling_selection = nodes.new("FunctionNodeBooleanMath") + culling_selection.operation = "OR" + culling_selection.location = (-550.0, 300.0) + links.new(final_or.outputs[0], culling_selection.inputs[0]) + links.new(full_stretch.outputs["Result"], culling_selection.inputs[1]) + delete_geometry = nodes.new("GeometryNodeDeleteGeometry") + delete_geometry.domain = "FACE" + delete_geometry.location = (-350.0, 100.0) + links.new(geometry_input.outputs["Geometry"], delete_geometry.inputs["Geometry"]) + links.new(culling_selection.outputs[0], delete_geometry.inputs["Selection"]) + + vertex_neighbors = nodes.new("GeometryNodeInputMeshVertexNeighbors") + vertex_neighbors.location = (-550.0, -320.0) + orphan_vertex = new_compare( + nodes, + "INT", + "EQUAL", + (-350.0, -320.0), + 0, + ) + links.new( + vertex_neighbors.outputs["Face Count"], + input_socket(orphan_vertex, "A", 0), + ) + delete_orphans = nodes.new("GeometryNodeDeleteGeometry") + delete_orphans.domain = "POINT" + delete_orphans.location = (-300.0, 100.0) + links.new(delete_geometry.outputs["Geometry"], delete_orphans.inputs["Geometry"]) + links.new(orphan_vertex.outputs["Result"], delete_orphans.inputs["Selection"]) + + remeshed = add_remesh_zone( + group, + delete_orphans.outputs["Geometry"], + remesh_input.outputs["Resolution"], + remesh_input.outputs["Relax Iterations"], + topology_input.outputs["Fix Topology"], + remesh_input.outputs["Smooth"], + (-300.0, 100.0), + ) + + mesh_island = nodes.new("GeometryNodeInputMeshIsland") + mesh_island.location = (3400.0, 700.0) + face_area = nodes.new("GeometryNodeInputMeshFaceArea") + face_area.location = (3400.0, 550.0) + island_area = nodes.new("GeometryNodeAccumulateField") + island_area.data_type = "FLOAT" + island_area.domain = "FACE" + island_area.location = (3600.0, 630.0) + links.new(face_area.outputs["Area"], island_area.inputs["Value"]) + links.new(mesh_island.outputs["Island Index"], island_area.inputs["Group ID"]) + resolution_squared = nodes.new("ShaderNodeMath") + resolution_squared.operation = "POWER" + resolution_squared.inputs[1].default_value = 2.0 + resolution_squared.location = (3600.0, 250.0) + links.new(island_input.outputs["Resolution"], resolution_squared.inputs[0]) + tiny_island = new_compare( + nodes, + "FLOAT", + "LESS_THAN", + (3800.0, 630.0), + 0.0, + ) + links.new(island_area.outputs["Total"], input_socket(tiny_island, "A", 0)) + links.new(resolution_squared.outputs[0], input_socket(tiny_island, "B", 1)) + remove_tiny = nodes.new("GeometryNodeDeleteGeometry") + remove_tiny.domain = "FACE" + remove_tiny.location = (4000.0, 100.0) + links.new(remeshed, remove_tiny.inputs["Geometry"]) + links.new(tiny_island.outputs["Result"], remove_tiny.inputs["Selection"]) + + reprojected = add_uv_reprojection( + group, + uv_source_input.outputs["Geometry"], + remove_tiny.outputs["Geometry"], + (4100.0, 100.0), + ) + shade_smooth = nodes.new("GeometryNodeSetShadeSmooth") + shade_smooth.domain = "FACE" + shade_smooth.location = (5400.0, 100.0) + shade_smooth.inputs["Shade Smooth"].default_value = True + links.new(reprojected, shade_smooth.inputs["Geometry"]) + links.new(shade_smooth.outputs["Geometry"], group_output.inputs["Geometry"]) + + group.color_tag = "GEOMETRY" + group.description = ( + "Cull DA3/BEN2 faces, repair merged topology, preserve boundaries, " + "and reproject UVs." + ) + group.use_fake_user = True + return group + + +def main(): + bpy.ops.wm.read_factory_settings(use_empty=True) + bpy.context.preferences.filepaths.save_version = 0 + build_group() + OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True) + bpy.ops.wm.save_as_mainfile(filepath=str(OUTPUT_PATH), check_existing=False) + + +if __name__ == "__main__": + main() diff --git a/tools/validate_geometry_nodes_asset.py b/tools/validate_geometry_nodes_asset.py new file mode 100644 index 0000000..7c2c976 --- /dev/null +++ b/tools/validate_geometry_nodes_asset.py @@ -0,0 +1,188 @@ +from pathlib import Path +import sys + +import bpy +import numpy as np + + +sys.path.insert(0, str(Path(__file__).parents[1] / "src")) + +from snapform.image_mesh import create_matching_camera, load_geometry_node_group + + +def set_input(modifier, group, name, value): + item = next( + item + for item in group.interface.items_tree + if item.item_type == "SOCKET" + and item.in_out == "INPUT" + and item.name == name + ) + modifier[item.identifier] = value + + +def add_face_attributes(mesh, count, background_values=None): + for name in ("o_face_confidence", "o_edge_stretch"): + attribute = mesh.attributes.new(name, "FLOAT", "FACE") + attribute.data.foreach_set("value", [1.0] * count) + background = mesh.attributes.new("o_background_face", "BOOLEAN", "FACE") + background.data.foreach_set( + "value", + background_values or [False] * count, + ) + + +def main(): + bpy.ops.wm.read_factory_settings(use_empty=True) + group = load_geometry_node_group() + assert load_geometry_node_group() == group + assert group.is_modifier + assert group.asset_data is None + assert not any(node.bl_idname == "GeometryNodeGroup" for node in group.nodes) + assert not any(node.bl_idname == "GeometryNodeTriangulate" for node in group.nodes) + group_inputs = [ + node for node in group.nodes if node.bl_idname == "NodeGroupInput" + ] + assert len(group_inputs) >= 8 + for node in group_inputs: + visible_outputs = [ + output for output in node.outputs if output.name and not output.hide + ] + assert 1 <= len(visible_outputs) <= 3 + uv_stores = [ + node + for node in group.nodes + if node.bl_idname == "GeometryNodeStoreNamedAttribute" + and node.inputs["Name"].default_value == "UVMap" + ] + assert len(uv_stores) == 1 + assert uv_stores[0].data_type == "FLOAT2" + panels = { + item.name: item + for item in group.interface.items_tree + if item.item_type == "PANEL" + } + assert not panels["Culling"].default_closed + assert panels["Options"].default_closed + assert [item.name for item in panels["Culling"].interface_items] == [ + "Confidence", + "Stretch", + "Background", + ] + assert [item.name for item in panels["Options"].interface_items] == [ + "Fix Topology", + ] + if bpy.app.version >= (5, 0, 0): + assert group.library is not None + + mesh = bpy.data.meshes.new("Validation Mesh") + mesh.from_pydata( + [(0, 0, 0), (1, 0, 0), (0, 1, 0), (1, 1, 0)], + [], + [(0, 2, 1), (1, 2, 3)], + ) + uv = mesh.uv_layers.new(name="UVMap") + uv.data.foreach_set("uv", [0, 0, 0, 1, 1, 0, 1, 0, 0, 1, 1, 1]) + add_face_attributes(mesh, 2, [False, True]) + + obj = bpy.data.objects.new("Validation Object", mesh) + bpy.context.scene.collection.objects.link(obj) + modifier = obj.modifiers.new(name=group.name, type="NODES") + modifier.node_group = group + modifier.show_group_selector = False + assert not modifier.show_group_selector + set_input(modifier, group, "Resolution", 0.000001) + set_input(modifier, group, "Relax Iterations", 1) + set_input(modifier, group, "Fix Topology", 0) + set_input(modifier, group, "Smooth", 0.0) + evaluated = obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) + result = evaluated.to_mesh() + try: + assert len(result.polygons) == 1 + assert "UVMap" in result.attributes + assert result.attributes["UVMap"].data_type == "FLOAT2" + finally: + evaluated.to_mesh_clear() + + set_input(modifier, group, "Stretch", 1.0) + stretch_item = next( + item + for item in group.interface.items_tree + if item.item_type == "SOCKET" and item.name == "Stretch" + ) + assert modifier[stretch_item.identifier] == 1.0 + obj.update_tag(refresh={"DATA"}) + bpy.context.view_layer.update() + culled = obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) + culled_result = culled.to_mesh() + try: + assert len(culled_result.polygons) == 0 + assert len(culled_result.vertices) == 0 + finally: + culled.to_mesh_clear() + + size = 10 + dense_vertices = [ + (x * 0.1, y * 0.1, 0.0) + for y in range(size) + for x in range(size) + ] + dense_faces = [] + for y in range(size - 1): + for x in range(size - 1): + a = y * size + x + b = a + 1 + c = a + size + d = c + 1 + dense_faces.extend(((a, c, b), (b, c, d))) + dense_mesh = bpy.data.meshes.new("Dense Validation Mesh") + dense_mesh.from_pydata(dense_vertices, [], dense_faces) + dense_uv = dense_mesh.uv_layers.new(name="UVMap") + dense_uv.data.foreach_set("uv", [0.0] * (len(dense_mesh.loops) * 2)) + add_face_attributes(dense_mesh, len(dense_faces)) + dense_object = bpy.data.objects.new("Dense Validation Object", dense_mesh) + bpy.context.scene.collection.objects.link(dense_object) + dense_modifier = dense_object.modifiers.new(name=group.name, type="NODES") + dense_modifier.node_group = group + set_input(dense_modifier, group, "Resolution", 0.22) + dense_evaluated = dense_object.evaluated_get( + bpy.context.evaluated_depsgraph_get() + ) + dense_result = dense_evaluated.to_mesh() + try: + assert 4 < len(dense_result.polygons) < len(dense_faces) + assert all(polygon.normal.z < -0.01 for polygon in dense_result.polygons) + finally: + dense_evaluated.to_mesh_clear() + + source = bpy.data.objects.new("Validation Source", None) + bpy.context.scene.collection.objects.link(source) + source.location = (2.0, 3.0, 4.0) + bpy.context.view_layer.update() + obj.matrix_world = source.matrix_world.copy() + bpy.context.view_layer.update() + camera = create_matching_camera( + bpy.context, + source, + obj, + "Validation", + np.asarray( + ((100.0, 0.0, 50.0), (0.0, 100.0, 50.0), (0.0, 0.0, 1.0)) + ), + (101, 101), + {"scale_x": 1.0, "scale_z": 1.0, "location": (0.0, 0.0, 2.0)}, + ) + assert camera.parent == obj + camera_location = tuple( + round(value, 5) for value in camera.matrix_world.translation + ) + assert camera_location == ( + 2.0, + 3.0, + 6.0, + ) + print(f"Validated Blender {bpy.app.version_string}: {group.name}") + + +if __name__ == "__main__": + main()