diff --git a/README.md b/README.md index d7ae0d8..d48c5bd 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ SnapForm 是一个 Blender 扩展，可以把一张照片、一段视频或一 5. 打开 3D 视图右侧的 `SnapForm` 面板，点击 `Install Environment` DA3、MoGe-2、utils3d、BEN2 源码包和 uv 已内置在扩展里，安装时直接从本地解包，不需要访问 GitHub。首次使用仍会自动 -下载运行环境依赖（PyTorch 等），随后自动下载 `DA3 Small` 和 `BEN2 Base`。其他 AI 模型可在 +下载运行环境依赖（PyTorch 等），随后自动下载 `MoGe-2 ViT-B Normal` 和 `BEN2 Base`。其他 AI 模型可在 `Models` panel 中按需下载，请保持网络畅通，首次安装需要一些时间。 ## 快速上手 @@ -44,14 +44,16 @@ DA3、MoGe-2、utils3d、BEN2 源码包和 uv 已内置在扩展里，安装时 ### 转换为透视网格 -选中单张图片的 Image Empty，右键打开 `SnapForm` 子菜单并选择 `Convert to Surface`。SnapForm 会按 DA3 -相机内参与深度逐像素反投影，并把生成网格匹配到原 Empty 的可见尺寸、Offset 与 Transform。DA3 -confidence 与边拉伸结果会保存为 Surface attributes，再由一个可复用的 Geometry Nodes +选中单张图片的 Image Empty，右键打开 `SnapForm` 子菜单并选择 `Convert to Surface`。可选择 MoGe-2 +ViT-S、ViT-B、ViT-L 或 DA3 Small、Base，其中 MoGe-2 ViT-B 为默认模型。SnapForm 会按模型输出的 +相机内参与三维点建立透视网格，并匹配原 Empty 的可见尺寸、Offset 与 Transform。MoGe-2 直接使用官方 +`utils3d.build_mesh_from_map` 与深度断层 mask 生成干净拓扑；DA3 保留 confidence 与边拉伸剔除。两者再由独立的 Geometry Nodes 修改器完成剔面、快速重拓扑、Relax 与 UV 重投。开启 Remove Background 时，背景三角面会在转换期间 直接删除，不再交给 Geometry Nodes。匹配相机默认不创建；开启后会作为 Surface 的子级。 当前此功能仅支持单张图片。 -转换对话框提供模型与 `Resolution Scale`、`Create Matching Camera`、`Remove Background`、`Unshaded`， +转换对话框提供模型与 `Create Matching Camera`、`Remove Background`、`Unshaded`；DA3 显示 +`Resolution Scale`，MoGe-2 显示 `Resolution Level`， 以及分割线下方的 `Keep Original`。 界面会显示计算所得的实际 Process Res。剔除和重拓扑参数直接在 Geometry Nodes 修改器中调整；`Unshaded` 可使用与 Depth Plane 相同的 Emission 材质。普通材质的 diff --git a/docs/technical.md b/docs/technical.md index 5f11c76..e0ffa9b 100644 --- a/docs/technical.md +++ b/docs/technical.md @@ -25,7 +25,7 @@ - Environment 已安装时不额外占用面板空间；模型列表每个模型只占一行 - `Debug` 默认折叠，内部的 `DA3 Output Test`、`MoGe-2 API Test` 与 `BEN2 Removal Test` 也分别默认折叠 - 使用一个输出类型下拉菜单和一个 `Generate Output` 按钮 -- 只显示当前格式的专属参数；GLB 与 PLY 设置不会同时出现 +- DA3 与 MoGe-2 Debug 都先选择单一 Output Format，只显示当前格式的专属参数；GLB 与 PLY 不会在一次测试中同时生成 - Debug 的模型选项按格式过滤；PLY 只提供两个带 Gaussian head 的 Giant 模型 - Debug 与 Convert to Depth Plane 的模型选项直接显示下载状态；未下载模型带 `(Missing)`，已下载模型保持原名 @@ -110,7 +110,7 @@ Cyclic 沿用原 Empty 的设置。序列帧无法被 Blender 整体 Pack，因 `Models` panel 下载所选模型，并固定使用 `process_res_method="upper_bound_resize"`；设备和参考视图策略 继续使用场景设置。Model 下拉菜单只列出已经下载的 DA3 模型；若没有可用模型，对话框会引导到 `Models` panel 下载。 - - `MoGe-2`：使用 Base Normal 模型输出 metric depth、normal、validity mask 与 normalized intrinsics。正式转换的 `results.npz` 不保存体积较大的 points；Debug 测试会保留 points，并另外写出 depth、normal、mask 预览。 + - `MoGe-2`：使用所选模型输出 metric depth、normal、validity mask 与 normalized intrinsics。Depth Plane 不保存体积较大的 points；Surface 会在 AI 进程内用 points 和官方 `utils3d.build_mesh_from_map` 直接生成 `mesh.npz`。Debug 始终保留原始 `results.npz`，并按 Output Format 单独生成 Maps、纹理 GLB 或彩色 PLY。 - `Original`：直接把原图作为深度贴图，不显示估算参数。仅在开启 `Remove Background` 时运行 BEN2 取得前景 Alpha；视频或序列帧按需同时抽取原始帧与前景序列，把原始帧作为 `.depth` 序列。 `Original` 是默认深度来源。选择 `DA3` 时，其 Model、Use Ray Pose 与 Resolution Scale @@ -168,8 +168,8 @@ Cyclic 沿用原 Empty 的设置。序列帧无法被 Blender 整体 Pack，因 1. 在 3D View 选中仅包含一帧的 Image Empty，右键打开 `SnapForm` 子菜单并选择 `Convert to Surface`。该入口独立于 `Convert to Depth Plane`，当前不处理视频或序列帧。 - 对话框中的 Resolution Scale 与 Depth Plane 使用相同设计：原图长边乘以 `1–100%` 的比例后， - 将实际 Process Res 钳制在 `252–1008` 并按 `14` 对齐。 + 模型列表包含三个 MoGe-2 Normal 模型和两个 DA3 模型，默认使用 MoGe-2 ViT-B Normal。DA3 的 + Resolution Scale 与 Depth Plane 使用相同设计；MoGe-2 改为显示共享的 Resolution Level。 2. Service 请求 DA3 的 `depth`、`conf` 与 `intrinsics`。仅在开启 `Remove Background` 时运行 BEN2 取得 前景 Alpha，并在输出前删除包含背景顶点的三角面；关闭时不加载 BEN2。Depth 不做 smoothing，Service 反投影深度网格，建立源三角面，计算 confidence percentile 与 camera-space edge stretch，并写入 @@ -189,6 +189,11 @@ Cyclic 沿用原 Empty 的设置。序列帧无法被 Blender 整体 Pack，因 - Smooth 采用 `O Smooth` 的边界检测和模糊法线权重，固定开放边界与薄面，只保留一个 `Smooth` 程度系数；节点组不再额外执行 Flip Faces； - 从修改前的原始表面以 nearest-surface 重新采样 `UVMap`，最终明确写为 Corner 域 `2D Vector`。 + MoGe-2 使用独立的 `O MoGe2 Surface`。Service 与官方 `infer.py` 一样，先把模型 mask 与 + `utils3d.np.depth_map_edge(depth, rtol=0.04)` 合并，再由 `utils3d.np.build_mesh_from_map(..., tri=True)` + 直接生成紧凑三角网格；Blender 不再从规则像素网格重新生面。节点组保留顶层 `Mask`、Resolution、 + Relax、Smooth、拓扑修复与 UV 重投影，不提供 Confidence、Stretch 或 Remesh 开关；Mask 的 Structure + Type 为 `Single`。Remesh 仍是节点处理流程的一部分。 4. 节点组随插件放在 `O_SnapForm.blend`。Blender 4.5 首次 Append 后复用同一 datablock； Blender 5.0 以 linked ID 载入后调用 Pack，后续 Surface 继续复用。插件不会为每个 Surface 新建节点组， 也不会在添加后重命名 Geometry Nodes 修改器。节点组只启用 `Usage > Modifier`，不 Mark as Asset； @@ -202,7 +207,8 @@ Cyclic 沿用原 Empty 的设置。序列帧无法被 Blender 整体 Pack，因 数据后直接创建 Surface 与原图材质，不需要 Open3D 或 OpenCV 进入 Blender Python 环境。 Surface 保留下列结构化 attributes： - 点域：`o_depth`、`o_confidence`、`o_confidence_percentile`、`o_background`； - - 面域：`o_face_confidence`、`o_edge_stretch`、`o_background_face`。 + - 面域：`o_face_confidence`、`o_background_face`；`o_edge_stretch` 位于边域。 + MoGe-2 的导入网格只保留点域 `o_depth` 与面域 `o_mask`，不伪造 confidence 或 background 属性。 7. 网格以 `reference_depth` 截面匹配 Image Empty 的显示矩形：该截面的投影四角与 Empty 的可见 边界重合，主点同时匹配 Empty Offset。网格对象直接沿用源 Empty 的 `matrix_world`，因此位置、旋转、 缩放及父级关系产生的世界变换保持一致。更近的顶点沿 Empty 局部 `+Z` 方向凸出，三角面正面朝向 @@ -268,7 +274,8 @@ Displacement 设置为 `Displacement Only`，DA3 与 Original 都使用 `Midleve `Background Depth`，使两种置换路径保持一致；未移除背景的普通不透明图片因此全部视为前景。 视频与序列帧的颜色、深度与前景 Image Texture 通过 image_user 的 Frames / Offset / Start 把第 i 帧映射到场景 `Start+i` 帧，并启用 Auto Refresh。 -颜色贴图到 Principled BSDF 之间使用独立的 Shader Node Group。Surface 使用 `O DA3 Surface Layer`； +颜色贴图到 Principled BSDF 之间使用独立的 Shader Node Group。DA3 Surface 使用 `O DA3 Surface Layer`， +MoGe-2 Surface 使用独立的 `O MoGe2 Surface Layer`； Depth Plane 使用 `O DA3 Depth Plane Layer` 材质组，并把背景剔除、背景深度与材质置换 Attribute 读取全部收进组内。两者都接收 `Color`、`Alpha` 与 `Bump Scale`，只有 `O DA3 Depth Plane Layer` 额外接收 Float 类型的 `Depth`。 diff --git a/src/snapform/assets/O_SnapForm.blend b/src/snapform/assets/O_SnapForm.blend index 4e6b0f4..2003356 100644 --- a/src/snapform/assets/O_SnapForm.blend +++ b/src/snapform/assets/O_SnapForm.blend @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:a762ec6c2dfb8e3a1d13bc3a937a2ce30d69195f91b05cde45697f38170b220d -size 1079655 +oid sha256:7ce30daf46cec0d4d6286cb2913b05a9a9f00a3fd9fe83de6f9e494e511ad8da +size 1314790 diff --git a/src/snapform/client.py b/src/snapform/client.py index 0dc3d98..a24f070 100644 --- a/src/snapform/client.py +++ b/src/snapform/client.py @@ -38,6 +38,7 @@ JOB_FOLDER_NAMES = { "RECONSTRUCT": "export-reconstruction", "IMAGE_DEPTH": "generate-depth-plane", "MOGE2_DEPTH": "generate-moge2-depth-plane", + "MOGE2_MESH": "generate-moge2-surface", "MOGE2_DEBUG": "test-moge2-api", "IMAGE_MESH": "generate-image-mesh", "REMOVE_BACKGROUND": "remove-background", @@ -909,6 +910,8 @@ def moge_request( start_frame=1, device=None, model_key=None, + debug_export_format="", + surface_mesh=False, ): from .models import BEN2_MODEL, MOGE2_MODEL @@ -929,6 +932,8 @@ def moge_request( else resolution_level ), "include_points": bool(include_points), + "debug_export_format": debug_export_format, + "surface_mesh": bool(surface_mesh), "video_frames": settings.video_frames, "start_frame": start_frame, "device": (device or settings.device).lower(), diff --git a/src/snapform/image_depth.py b/src/snapform/image_depth.py index 1d9215a..fd45373 100644 --- a/src/snapform/image_depth.py +++ b/src/snapform/image_depth.py @@ -19,6 +19,7 @@ MATERIAL_DISPLACEMENT_ATTRIBUTE = "o_material_displacement" BACKGROUND_DEPTH_ATTRIBUTE = "o_background_depth" DEPTH_SCALE_ATTRIBUTE = "o_depth_scale" SURFACE_MATERIAL_NODE_GROUP_NAME = "O DA3 Surface Layer" +MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME = "O MoGe2 Surface Layer" DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME = "O DA3 Depth Plane Layer" MOGE_MATERIAL_NODE_GROUP_NAME = "O MoGe2 Depth Plane Layer" DA3_PATCH_SIZE = 14 @@ -57,9 +58,6 @@ def image_depth_model_items(owner, context): return depth_plane_model_enum_items(owner, context) -_last_depth_plane_model = "MOGE2_VITB_NORMAL" - - def is_image_empty(candidate): return bool( candidate @@ -1006,10 +1004,13 @@ def create_moge_material( uv.location = (-900.0, 200.0) color = nodes.new("ShaderNodeTexImage") color.image = color_image + color.interpolation = "Linear" + color.extension = "EXTEND" color.location = (-650.0, 300.0) depth = nodes.new("ShaderNodeTexImage") depth.image = depth_image depth.interpolation = "Linear" + depth.extension = "EXTEND" depth.location = (-650.0, -100.0) if frame_count > 1: configure_image_user( @@ -1127,11 +1128,11 @@ def create_plane_mesh(source_object, bounds=None, subdivision_level=5): return mesh -def material_node_group(depth_plane=False): +def material_node_group(depth_plane=False, surface_group_name=None): group_name = ( DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME if depth_plane - else SURFACE_MATERIAL_NODE_GROUP_NAME + else surface_group_name or SURFACE_MATERIAL_NODE_GROUP_NAME ) existing = bpy.data.node_groups.get(group_name) if existing is not None: @@ -1269,7 +1270,7 @@ def material_node_group(depth_plane=False): group.description = ( "Prepare DA3 depth-plane material outputs." if depth_plane - else "Prepare DA3 surface material outputs." + else f"Prepare {group_name} material outputs." ) group.use_fake_user = True return group @@ -1293,6 +1294,7 @@ def create_material( depth_frame_offset=0, depth_frame_duration=None, use_cyclic=None, + surface_group_name=None, ): material = bpy.data.materials.new(name) material.use_nodes = True @@ -1347,7 +1349,10 @@ def create_material( alpha = color_texture.outputs["Alpha"] material_inputs = nodes.new("ShaderNodeGroup") is_depth_plane = displacement_image is not None - material_inputs.node_tree = material_node_group(is_depth_plane) + material_inputs.node_tree = material_node_group( + is_depth_plane, + surface_group_name=surface_group_name, + ) material_inputs.name = material_inputs.node_tree.name material_inputs.label = material_inputs.name material_inputs.location = (-300.0, 350.0) @@ -2232,35 +2237,21 @@ class ConvertImageToDepthPlane(bpy.types.Operator): ) def invoke(self, context, _event): - global _last_depth_plane_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_depth_plane_model - if _last_depth_plane_model in available_keys - else available_keys[0] - ) return context.window_manager.invoke_props_dialog(self, width=340) def draw(self, context): layout = self.layout layout.prop(self, "snapform_model") if self.snapform_model in {"SMALL", "BASE"}: - estimate_box = layout.box() - estimate_box.prop(self, "snapform_use_ray_pose") - draw_process_resolution(estimate_box, self, context.object.data) + layout.prop(self, "snapform_use_ray_pose") + draw_process_resolution(layout, self, context.object.data) elif self.snapform_model.startswith("MOGE2_"): from .client import model_ready - moge_box = layout.box() if not model_ready(self.snapform_model): - moge_box.label(text="MoGe-2 model is not installed", icon="ERROR") - moge_box.label(text="Download it in Models") - moge_box.prop(self, "snapform_moge_resolution_level") + layout.label(text="MoGe-2 model is not installed", icon="ERROR") + layout.label(text="Download it in Models") + layout.prop(self, "snapform_moge_resolution_level") layout.prop(self, "snapform_remove_background") layout.prop(self, "snapform_unshaded") layout.separator(type="LINE") @@ -2369,8 +2360,6 @@ class ConvertImageToDepthPlane(bpy.types.Operator): if delete_input: remove_temporary_input(source_path) return {"CANCELLED"} - global _last_depth_plane_model - _last_depth_plane_model = self.snapform_model return {"FINISHED"} diff --git a/src/snapform/image_mesh.py b/src/snapform/image_mesh.py index db4314b..8eec143 100644 --- a/src/snapform/image_mesh.py +++ b/src/snapform/image_mesh.py @@ -3,9 +3,9 @@ from pathlib import Path import bpy from .image_depth import ( + MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME, create_material, draw_process_resolution, - image_depth_model_items, is_image_empty, plane_bounds, prepare_image_input, @@ -16,8 +16,6 @@ from .image_depth import ( ) -_last_mesh_model = "BASE" - POINT_ATTRIBUTES = { "o_depth": ("FLOAT", "point_depth"), "o_confidence": ("FLOAT", "point_confidence"), @@ -32,19 +30,28 @@ FACE_ATTRIBUTES = { "FLOAT", "face_confidence_percentile", ), - "o_edge_stretch": ("FLOAT", "face_edge_stretch"), "o_background_face": ("BOOLEAN", "face_background"), } +OPTIONAL_FACE_ATTRIBUTES = { + "o_mask": ("BOOLEAN", "face_mask"), +} NODE_GROUP_NAME = "O DA3 Surface" +MOGE_NODE_GROUP_NAME = "O MoGe2 Surface" NODE_ASSET_NAME = "O_SnapForm.blend" -def load_image_mesh_data(npz_path): +def image_surface_model_items(owner, context): + from .models import surface_model_enum_items + + return surface_model_enum_items(owner, context) + + +def load_image_mesh_data(npz_path, moge2=False): import numpy as np try: with np.load(npz_path, allow_pickle=False) as archive: - required = ( + required = [ "vertices", "faces", "uv", @@ -52,13 +59,13 @@ def load_image_mesh_data(npz_path): "image_size", "reference_depth", "point_depth", - "point_confidence", - "point_confidence_percentile", - "point_background", - "face_confidence_percentile", - "face_edge_stretch", - "face_background", - ) + ] + if moge2: + required.append("face_mask") + else: + required.extend( + ("point_confidence", "point_confidence_percentile", "point_background", "face_confidence_percentile", "face_background") + ) missing = [key for key in required if key not in archive.files] if missing: raise ValueError( @@ -81,7 +88,16 @@ def load_image_mesh_data(npz_path): **POINT_ATTRIBUTES, **FACE_ATTRIBUTES, }.items() + if key in archive.files } + for _name, (_data_type, key) in OPTIONAL_FACE_ATTRIBUTES.items(): + if key in archive.files: + attributes[key] = np.asarray(archive[key]) + if "vertex_normal" in archive.files: + attributes["vertex_normal"] = np.asarray( + archive["vertex_normal"], + dtype=np.float32, + ) except OSError as error: raise ValueError(f"Unable to read surface data: {npz_path}") from error if vertices.ndim != 2 or vertices.shape[1] != 3: @@ -97,10 +113,13 @@ def load_image_mesh_data(npz_path): if not reference_depth > 0.0: raise ValueError("Surface reference depth must be positive") for _name, (_data_type, key) in POINT_ATTRIBUTES.items(): - if len(attributes[key]) != len(vertices): + if key in attributes and 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): + if key in attributes and len(attributes[key]) != len(faces): + raise ValueError(f"Face attribute has the wrong length: {key}") + for _name, (_data_type, key) in OPTIONAL_FACE_ATTRIBUTES.items(): + if key in attributes and len(attributes[key]) != len(faces): raise ValueError(f"Face attribute has the wrong length: {key}") return ( vertices, @@ -185,7 +204,11 @@ def assign_mesh_attribute(mesh, name, data_type, domain, values): attribute.data.foreach_set("value", array.reshape(-1)) -def assign_snapform_attributes(mesh, attributes): +def assign_snapform_attributes(mesh, attributes, moge2=False): + if moge2: + assign_mesh_attribute(mesh, "o_depth", "FLOAT", "POINT", attributes["point_depth"]) + assign_mesh_attribute(mesh, "o_mask", "BOOLEAN", "FACE", attributes["face_mask"]) + return for name, (data_type, key) in POINT_ATTRIBUTES.items(): assign_mesh_attribute( mesh, @@ -202,6 +225,36 @@ def assign_snapform_attributes(mesh, attributes): "FACE", attributes[key], ) + for name, (data_type, key) in OPTIONAL_FACE_ATTRIBUTES.items(): + if key in attributes: + assign_mesh_attribute( + mesh, + name, + data_type, + "FACE", + attributes[key], + ) + assign_edge_depth_difference(mesh, attributes["point_depth"]) + + +def assign_edge_depth_difference(mesh, point_depth): + import numpy as np + + edge_vertices = np.empty(len(mesh.edges) * 2, dtype=np.int32) + mesh.edges.foreach_get("vertices", edge_vertices) + edge_vertices = edge_vertices.reshape(-1, 2) + depth = np.asarray(point_depth, dtype=np.float32) + first = depth[edge_vertices[:, 0]] + second = depth[edge_vertices[:, 1]] + denominator = np.maximum(np.minimum(np.abs(first), np.abs(second)), 1e-6) + difference = np.abs(first - second) / denominator + assign_mesh_attribute( + mesh, + "o_edge_stretch", + "FLOAT", + "EDGE", + difference, + ) def remove_surface_background(vertices, faces, uv, attributes): @@ -228,12 +281,100 @@ def remove_surface_background(vertices, faces, uv, attributes): ) +def prepare_moge_surface_archive(results_path, output_path): + import numpy as np + + with np.load(results_path, allow_pickle=False) as archive: + required = {"depth", "mask", "points", "intrinsics"} + missing = required.difference(archive.files) + if missing: + raise ValueError(f"MoGe-2 surface data is missing: {', '.join(sorted(missing))}") + depth = np.asarray(archive["depth"], dtype=np.float32) + mask = np.asarray(archive["mask"], dtype=bool) + points = np.asarray(archive["points"], dtype=np.float32) + intrinsics = np.asarray(archive["intrinsics"], dtype=np.float32) + normal = ( + np.asarray(archive["normal"], dtype=np.float32) + if "normal" in archive.files + else None + ) + if depth.shape[0] != 1 or points.shape[:3] != (*depth.shape,): + raise ValueError("Convert to Surface requires one matching MoGe-2 point grid") + raw_depth = depth[0] + finite = np.isfinite(raw_depth) & (raw_depth > 0.0) + finite &= np.isfinite(points[0]).all(axis=-1) + height, width = raw_depth.shape + indices = np.arange(height * width, dtype=np.int32).reshape(height, width) + quads = np.stack( + (indices[:-1, :-1], indices[1:, :-1], indices[1:, 1:], indices[:-1, 1:]), + axis=-1, + ).reshape(-1, 4) + faces = np.concatenate((quads[:, (0, 1, 2)], quads[:, (0, 2, 3)]), axis=0) + faces = faces[finite.reshape(-1)[faces].all(axis=1)] + if not len(faces): + raise ValueError("MoGe-2 output does not contain a valid surface") + flat_points = points[0].reshape(-1, 3) + clean_mask = mask[0] & ~depth_discontinuity_mask(raw_depth, 0.04) + face_mask = clean_mask.reshape(-1)[faces].all(axis=1) + used = np.unique(faces) + index_map = np.full(height * width, -1, dtype=np.int32) + index_map[used] = np.arange(len(used), dtype=np.int32) + u = np.linspace(0.0, 1.0, width, dtype=np.float32) + v = np.linspace(1.0, 0.0, height, dtype=np.float32) + uv = np.stack(np.meshgrid(u, v), axis=-1).reshape(-1, 2)[used] + camera_intrinsics = intrinsics[0].copy() + principal_x = float(camera_intrinsics[0, 2]) + principal_y = float(camera_intrinsics[1, 2]) + if 0.0 <= principal_x <= 1.0 and 0.0 <= principal_y <= 1.0: + camera_intrinsics[0] *= width + camera_intrinsics[1] *= height + camera_intrinsics[2] = (0.0, 0.0, 1.0) + point_depth = raw_depth.reshape(-1)[used] + masked_depth = raw_depth[finite & mask[0]] + reference_depth = float(np.median(masked_depth if len(masked_depth) else point_depth)) + np.savez( + output_path, + vertices=flat_points[used], + faces=index_map[faces], + uv=uv, + intrinsics=camera_intrinsics, + image_size=np.array((width, height), dtype=np.int32), + reference_depth=np.float32(reference_depth), + point_depth=point_depth, + vertex_normal=( + normal[0].reshape(-1, 3)[used] + if normal is not None + else np.empty((0, 3), dtype=np.float32) + ), + face_mask=face_mask, + ) + return output_path + + +def depth_discontinuity_mask(depth, threshold): + import numpy as np + + edges = np.zeros(depth.shape, dtype=bool) + for axis in (0, 1): + first = np.take(depth, range(depth.shape[axis] - 1), axis=axis) + second = np.take(depth, range(1, depth.shape[axis]), axis=axis) + scale = np.maximum(np.minimum(np.abs(first), np.abs(second)), 1e-6) + discontinuity = 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 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) +def load_geometry_node_group(group_name=NODE_GROUP_NAME): + existing = bpy.data.node_groups.get(group_name) if existing is not None: return existing asset_path = node_asset_path() @@ -243,15 +384,15 @@ def load_geometry_node_group(): if not use_pack: directory = str(asset_path) + "/NodeTree/" result = bpy.ops.wm.append( - filepath=directory + NODE_GROUP_NAME, + filepath=directory + group_name, directory=directory, - filename=NODE_GROUP_NAME, + filename=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) + node_group = bpy.data.node_groups.get(group_name) if node_group is None: raise RuntimeError("Appended Geometry Nodes asset is unavailable") node_group.use_fake_user = True @@ -260,11 +401,11 @@ def load_geometry_node_group(): source, target, ): - if NODE_GROUP_NAME not in source.node_groups: + if group_name not in source.node_groups: raise RuntimeError( - f"Geometry Nodes asset is missing {NODE_GROUP_NAME}" + f"Geometry Nodes asset is missing {group_name}" ) - target.node_groups = [NODE_GROUP_NAME] + target.node_groups = [group_name] node_group = target.node_groups[0] if use_pack: linked_group = node_group @@ -307,8 +448,13 @@ def material_depth_scale(mesh_object, camera_parameters): ) -def add_geometry_nodes_processing(mesh_object, vertices, faces): - node_group = load_geometry_node_group() +def add_geometry_nodes_processing( + mesh_object, + vertices, + faces, + group_name=NODE_GROUP_NAME, +): + node_group = load_geometry_node_group(group_name) modifier = mesh_object.modifiers.new(name=node_group.name, type="NODES") modifier.node_group = node_group modifier.show_group_selector = False @@ -360,7 +506,10 @@ def build_image_mesh( create_camera, unshaded, remove_background=False, + moge2=False, ): + import numpy as np + ( camera_vertices, faces, @@ -369,7 +518,7 @@ def build_image_mesh( image_size, reference_depth, attributes, - ) = load_image_mesh_data(npz_path) + ) = load_image_mesh_data(npz_path, moge2=moge2) if remove_background: camera_vertices, faces, uv, attributes = remove_surface_background( camera_vertices, @@ -388,10 +537,24 @@ def build_image_mesh( mesh = bpy.data.meshes.new(source_name) mesh.from_pydata(vertices.tolist(), [], faces.tolist()) assign_vertex_uv(mesh, uv) - assign_snapform_attributes(mesh, attributes) + assign_snapform_attributes(mesh, attributes, moge2=moge2) mesh.update() for polygon in mesh.polygons: polygon.use_smooth = True + if moge2 and len(attributes.get("vertex_normal", ())) == len(vertices): + normals = attributes["vertex_normal"].copy() + normals[:, 0] /= camera_parameters["scale_x"] + normals[:, 1] /= -camera_parameters["scale_y"] + normals[:, 2] /= -camera_parameters["scale_z"] + lengths = np.linalg.norm(normals, axis=1, keepdims=True) + normals /= np.maximum(lengths, 1e-6) + set_custom_normals = getattr( + mesh, + "normals_split_custom_set_from_vertices", + None, + ) + if callable(set_custom_normals): + set_custom_normals(normals.tolist()) mesh_object = bpy.data.objects.new(source_name, mesh) context.collection.objects.link(mesh_object) mesh_object.matrix_world = source_object.matrix_world.copy() @@ -401,11 +564,19 @@ def build_image_mesh( source_object.data, depth_scale, unshaded=unshaded, + surface_group_name=( + MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME if moge2 else None + ), ) 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) + add_geometry_nodes_processing( + mesh_object, + vertices, + faces, + MOGE_NODE_GROUP_NAME if moge2 else NODE_GROUP_NAME, + ) if create_camera: create_matching_camera( context, @@ -453,6 +624,10 @@ class BuildImageMesh(bpy.types.Operator): default=False, options={"HIDDEN", "SKIP_SAVE"}, ) + snapform_moge2: bpy.props.BoolProperty( + default=False, + options={"HIDDEN", "SKIP_SAVE"}, + ) def execute(self, context): source_object = bpy.data.objects.get(self.snapform_source_object_name) @@ -468,6 +643,7 @@ class BuildImageMesh(bpy.types.Operator): self.snapform_create_camera, self.snapform_unshaded, self.snapform_remove_background, + self.snapform_moge2, ) except (OSError, RuntimeError, ValueError) as error: self.report({"ERROR"}, str(error)) @@ -484,7 +660,8 @@ class ConvertImageToMesh(bpy.types.Operator): snapform_model: bpy.props.EnumProperty( name="Model", description="Model used to estimate image depth", - items=image_depth_model_items, + items=image_surface_model_items, + default=1, ) snapform_use_ray_pose: bpy.props.BoolProperty( name="Use Ray Pose", @@ -492,6 +669,13 @@ class ConvertImageToMesh(bpy.types.Operator): default=True, ) snapform_process_res_factor: process_resolution_scale_property() + snapform_moge_resolution_level: bpy.props.IntProperty( + name="Resolution Level", + description="MoGe-2 inference resolution level; higher is finer and slower", + default=5, + min=0, + max=9, + ) snapform_create_camera: bpy.props.BoolProperty( name="Create Matching Camera", description="Create a perspective camera matching the source projection", @@ -515,7 +699,7 @@ class ConvertImageToMesh(bpy.types.Operator): @classmethod def poll(cls, context): - from .models import MODELS + from .models import MODELS, MOGE2_MODELS from .operators import active_job from .client import service_busy, model_ready @@ -524,29 +708,23 @@ class ConvertImageToMesh(bpy.types.Operator): and video_total_frames(context.object.data) == 1 and active_job() is None and not service_busy() - and any(model_ready(model_key) for model_key in MODELS) + and any( + model_ready(model_key) + for model_key in (*MOGE2_MODELS, *MODELS) + ) ) 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") - draw_process_resolution(layout, self, context.object.data) + if self.snapform_model.startswith("MOGE2_"): + layout.prop(self, "snapform_moge_resolution_level") + else: + layout.prop(self, "snapform_use_ray_pose") + draw_process_resolution(layout, self, context.object.data) layout.prop(self, "snapform_create_camera") layout.prop(self, "snapform_remove_background") layout.prop(self, "snapform_unshaded") @@ -554,17 +732,19 @@ class ConvertImageToMesh(bpy.types.Operator): layout.prop(self, "snapform_keep_original") def execute(self, context): - from .operators import IMAGE_MESH_JOB + from .operators import IMAGE_MESH_JOB, MOGE2_MESH_JOB source_object = context.object if video_total_frames(source_object.data) != 1: self.report({"ERROR"}, "Convert to Surface currently supports still images") return {"CANCELLED"} try: - process_res = scaled_process_resolution( - source_object.data, - self.snapform_process_res_factor, - ) + process_res = 1 + if not self.snapform_model.startswith("MOGE2_"): + 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, @@ -573,12 +753,17 @@ class ConvertImageToMesh(bpy.types.Operator): self.report({"ERROR"}, str(error)) return {"CANCELLED"} result = bpy.ops.snapform.run_job( - job_type=IMAGE_MESH_JOB, + job_type=( + MOGE2_MESH_JOB + if self.snapform_model.startswith("MOGE2_") + else 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_moge_resolution_level=self.snapform_moge_resolution_level, snapform_delete_source=not self.snapform_keep_original, snapform_delete_input=delete_input, snapform_mesh_create_camera=self.snapform_create_camera, @@ -589,6 +774,4 @@ class ConvertImageToMesh(bpy.types.Operator): 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/models.py b/src/snapform/models.py index 35ed118..d9a4a45 100644 --- a/src/snapform/models.py +++ b/src/snapform/models.py @@ -139,6 +139,8 @@ MODEL_NUMBERS = { _ENUM_ITEMS_CACHE = {} +_DEPTH_PLANE_ENUM_ITEMS_CACHE = {} +_SURFACE_ENUM_ITEMS_CACHE = {} def enum_items(supports_gs=None, show_download_status=False): @@ -213,9 +215,13 @@ def depth_plane_model_enum_items(_owner, _context): *MOGE2_MODELS.values(), *(model for model in MODELS.values() if model.enabled), ] + ready = tuple(model_ready(model.key) for model in models) + cached = _DEPTH_PLANE_ENUM_ITEMS_CACHE.get(ready) + if cached is not None: + return cached items = [] for index, model in enumerate(models): - label = model.label if model_ready(model.key) else f"{model.label} (Missing)" + label = model.label if ready[index] else f"{model.label} (Missing)" items.append( ( model.key, @@ -225,6 +231,22 @@ def depth_plane_model_enum_items(_owner, _context): ) ) items.append(("ORIGINAL", "Original", "Use the original image as depth", len(items))) + _DEPTH_PLANE_ENUM_ITEMS_CACHE[ready] = items + return items + + +def surface_model_enum_items(owner, context): + depth_items = depth_plane_model_enum_items(owner, context) + cache_key = id(depth_items) + cached = _SURFACE_ENUM_ITEMS_CACHE.get(cache_key) + if cached is not None: + return cached + items = [ + item + for item in depth_items + if item[0] != "ORIGINAL" + ] + _SURFACE_ENUM_ITEMS_CACHE[cache_key] = items return items diff --git a/src/snapform/operators.py b/src/snapform/operators.py index 4609a49..effd52d 100644 --- a/src/snapform/operators.py +++ b/src/snapform/operators.py @@ -20,6 +20,7 @@ from .models import ( RECONSTRUCTION_JOB = "RECONSTRUCT" IMAGE_DEPTH_JOB = "IMAGE_DEPTH" MOGE2_DEPTH_JOB = "MOGE2_DEPTH" +MOGE2_MESH_JOB = "MOGE2_MESH" MOGE2_DEBUG_JOB = "MOGE2_DEBUG" IMAGE_MESH_JOB = "IMAGE_MESH" BACKGROUND_REMOVAL_JOB = "REMOVE_BACKGROUND" @@ -31,6 +32,7 @@ IMAGE_EMPTY_JOBS = frozenset( { *DA3_IMAGE_JOBS, MOGE2_DEPTH_JOB, + MOGE2_MESH_JOB, IMAGE_BACKGROUND_JOB, ORIGINAL_BACKGROUND_JOB, ORIGINAL_FRAMES_JOB, @@ -40,6 +42,7 @@ PRODUCTION_AI_JOBS = frozenset( { *DA3_IMAGE_JOBS, MOGE2_DEPTH_JOB, + MOGE2_MESH_JOB, IMAGE_BACKGROUND_JOB, ORIGINAL_BACKGROUND_JOB, } @@ -163,6 +166,7 @@ class RunJob(bpy.types.Operator): (RECONSTRUCTION_JOB, "Export DA3 Result", ""), (IMAGE_DEPTH_JOB, "Convert Image to Depth Plane", ""), (MOGE2_DEPTH_JOB, "Convert Image with MoGe-2", ""), + (MOGE2_MESH_JOB, "Convert Image to Surface with MoGe-2", ""), (MOGE2_DEBUG_JOB, "Test MoGe-2 API", ""), (IMAGE_MESH_JOB, "Convert Image to Surface", ""), (BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""), @@ -250,7 +254,7 @@ class RunJob(bpy.types.Operator): export_format = "mini_npz" elif self.job_type == IMAGE_MESH_JOB: export_format = "mesh_data" - elif self.job_type in {MOGE2_DEPTH_JOB, MOGE2_DEBUG_JOB}: + elif self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB, MOGE2_DEBUG_JOB}: export_format = "moge_npz" job = ActiveJob( @@ -425,8 +429,12 @@ class RunJob(bpy.types.Operator): ), paths, ) - if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_DEBUG_JOB}: - model_key = self.snapform_model_key or MOGE2_MODEL.key + if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB, MOGE2_DEBUG_JOB}: + model_key = ( + settings.moge_model + if self.job_type == MOGE2_DEBUG_JOB + else self.snapform_model_key or MOGE2_MODEL.key + ) model = get_downloadable_model(model_key) if not client.model_ready(model_key): raise RuntimeError(f"Download {model.label} first") @@ -440,17 +448,23 @@ class RunJob(bpy.types.Operator): input_path=input_path, resolution_level=( self.snapform_moge_resolution_level - if self.job_type == MOGE2_DEPTH_JOB + if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB} else settings.moge_resolution_level ), - include_points=self.job_type == MOGE2_DEBUG_JOB, + include_points=self.job_type in {MOGE2_MESH_JOB, MOGE2_DEBUG_JOB}, remove_background=( - self.job_type == MOGE2_DEPTH_JOB + self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB} and remove_background ), start_frame=start_frame, device=self._task_device(settings), model_key=model_key, + debug_export_format=( + settings.moge_export_format + if self.job_type == MOGE2_DEBUG_JOB + else "" + ), + surface_mesh=self.job_type == MOGE2_MESH_JOB, ), paths, ) @@ -531,7 +545,7 @@ 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 in {MOGE2_DEPTH_JOB, MOGE2_DEBUG_JOB}: + if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB, MOGE2_DEBUG_JOB}: return "Starting MoGe-2 inference" if self.job_type == IMAGE_MESH_JOB: return "Starting perspective surface generation" @@ -728,6 +742,21 @@ class RunJob(bpy.types.Operator): self._redraw_ui(context, force=True) return {"CANCELLED"} status["message"] = "MoGe-2 depth plane is ready" + if job.kind == MOGE2_MESH_JOB: + try: + settings.status = "Building MoGe-2 surface" + settings.progress = 0.96 + self._redraw_ui(context, force=True) + self._build_moge2_image_mesh(job) + except (OSError, RuntimeError, ValueError) as import_error: + settings.status = "MoGe-2 surface import failed" + settings.progress = 0.0 + self.report({"ERROR"}, str(import_error)) + client.mark_service_job_complete(client.job_id_from_status(job.status_path)) + close_active_job() + self._redraw_ui(context, force=True) + return {"CANCELLED"} + status["message"] = "MoGe-2 surface is ready" if job.kind == IMAGE_MESH_JOB: try: settings.status = "Building perspective surface" @@ -801,6 +830,26 @@ class RunJob(bpy.types.Operator): if "FINISHED" not in undo_result: raise RuntimeError("Unable to create the surface Undo step") + def _build_moge2_image_mesh(self, job): + result_directory = job.result_directory + mesh_path = result_directory / "mesh.npz" + if not mesh_path.is_file(): + raise RuntimeError("MoGe-2 did not export the official surface mesh") + result = bpy.ops.snapform.build_image_mesh( + snapform_source_object_name=job.source_object_name, + snapform_npz_path=str(mesh_path), + snapform_delete_source=job.delete_source, + snapform_create_camera=job.mesh_create_camera, + snapform_unshaded=job.unshaded, + snapform_remove_background=False, + snapform_moge2=True, + ) + if "FINISHED" not in result: + raise RuntimeError("Unable to build the MoGe-2 perspective surface") + undo_result = bpy.ops.ed.undo_push(message="Convert to Surface") + if "FINISHED" not in undo_result: + raise RuntimeError("Unable to create the surface 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/panel.py b/src/snapform/panel.py index 3a2de93..b066817 100644 --- a/src/snapform/panel.py +++ b/src/snapform/panel.py @@ -252,9 +252,10 @@ class Moge2ApiDebugPanel(bpy.types.Panel): def draw(self, context): layout = self.layout settings = context.scene.snapform_settings + layout.prop(settings, "moge_export_format") layout.prop(settings, "input_path") - layout.prop(settings, "moge_resolution_level") layout.prop(settings, "moge_model") + layout.prop(settings, "moge_resolution_level") run_row = layout.row() run_row.enabled = ( active_job() is None diff --git a/src/snapform/properties.py b/src/snapform/properties.py index 9a03aa4..c093103 100644 --- a/src/snapform/properties.py +++ b/src/snapform/properties.py @@ -78,6 +78,16 @@ class SnapFormSettings(bpy.types.PropertyGroup): items=moge2_model_enum_items, default=1, ) + moge_export_format: bpy.props.EnumProperty( + name="Output Format", + description="Additional MoGe-2 debug output generated with results.npz", + items=( + ("maps", "Maps", "Export depth, normal, and mask previews"), + ("glb", "Textured Surface (GLB)", "Export the official-style textured mesh"), + ("ply", "Point Cloud (PLY)", "Export the official-style colored point cloud"), + ), + default="maps", + ) conf_thresh_percentile: bpy.props.FloatProperty( name="Confidence Threshold", description="Remove the lowest-confidence percentile of points or surface faces", diff --git a/src/snapform/worker.py b/src/snapform/worker.py index 009e024..db70a8e 100644 --- a/src/snapform/worker.py +++ b/src/snapform/worker.py @@ -1514,6 +1514,94 @@ def _save_moge_preview(output_directory, index, depth, normal, mask): ) +def export_moge_debug_mesh( + output_directory, + points, + depth, + mask, + normal, + image, + export_format, +): + import numpy as np + import utils3d + from moge.utils.io import save_glb, save_ply + + clean_mask = mask & ~utils3d.np.depth_map_edge(depth, rtol=0.04) + uv = utils3d.np.uv_map(*depth.shape) + result = utils3d.np.build_mesh_from_map( + points, + image, + uv, + normal, + mask=clean_mask, + tri=True, + ) + faces, vertices, colors, vertex_uvs, vertex_normals = result + vertices = vertices * np.array([1.0, -1.0, -1.0], dtype=np.float32) + vertex_uvs = vertex_uvs * np.array([1.0, -1.0]) + np.array([0.0, 1.0]) + vertex_normals = vertex_normals * np.array([1.0, -1.0, -1.0], dtype=np.float32) + if export_format == "glb": + save_glb( + output_directory / "mesh.glb", + vertices, + faces, + vertex_uvs, + (image * 255.0).astype(np.uint8), + vertex_normals, + ) + return + if export_format == "ply": + save_ply( + output_directory / "pointcloud.ply", + vertices, + np.zeros((0, 3), dtype=np.int32), + colors, + vertex_normals, + ) + + +def export_moge_surface_archive( + output_path, + points, + depth, + mask, + normal, + intrinsics, +): + import numpy as np + import utils3d + + height, width = depth.shape + clean_mask = mask & ~utils3d.np.depth_map_edge(depth, rtol=0.04) + result = utils3d.np.build_mesh_from_map( + points, + utils3d.np.uv_map(height, width), + normal, + mask=clean_mask, + tri=True, + ) + faces, vertices, vertex_uvs, vertex_normals = result + vertex_uvs = vertex_uvs * np.array([1.0, -1.0]) + np.array([0.0, 1.0]) + pixel_intrinsics = intrinsics.copy() + pixel_intrinsics[0] *= width + pixel_intrinsics[1] *= height + pixel_intrinsics[2] = (0.0, 0.0, 1.0) + point_depth = vertices[:, 2].astype(np.float32) + np.savez( + output_path, + vertices=vertices.astype(np.float32), + faces=faces.astype(np.int32), + uv=vertex_uvs.astype(np.float32), + intrinsics=pixel_intrinsics.astype(np.float32), + image_size=np.array((width, height), dtype=np.int32), + reference_depth=np.float32(np.median(point_depth)), + point_depth=point_depth, + vertex_normal=vertex_normals.astype(np.float32), + face_mask=np.ones(len(faces), dtype=bool), + ) + + def run_moge(args, model_cache=None, cancel_check=None): import numpy as np from PIL import Image @@ -1550,6 +1638,7 @@ def run_moge(args, model_cache=None, cancel_check=None): masks = [] points = [] intrinsics = [] + source_images = [] for index, image_path in enumerate(images): check_cancelled(cancel_check) write_status( @@ -1560,6 +1649,7 @@ def run_moge(args, model_cache=None, cancel_check=None): with Image.open(image_path) as opened: image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0 image_tensor = torch.from_numpy(image).permute(2, 0, 1).to(device) + source_images.append(image) with torch.inference_mode(): prediction = model.infer( image_tensor, @@ -1579,13 +1669,15 @@ def run_moge(args, model_cache=None, cancel_check=None): intrinsics.append(intrinsic) if bool(getattr(args, "include_points", False)): points.append(point) - _save_moge_preview( - output_directory, - None if len(images) == 1 else index, - depth, - normal, - mask, - ) + debug_format = getattr(args, "debug_export_format", "") + if not debug_format or debug_format == "maps": + _save_moge_preview( + output_directory, + None if len(images) == 1 else index, + depth, + normal, + mask, + ) shapes = {array.shape for array in depths} if len(shapes) != 1: @@ -1598,7 +1690,18 @@ def run_moge(args, model_cache=None, cancel_check=None): } if points: result["points"] = np.stack(points).astype(np.float32) - np.savez(output_directory / "results.npz", **result) + debug_format = getattr(args, "debug_export_format", "") + if debug_format in {"glb", "ply"} and len(images) == 1: + write_status(args.status, 0.84, f"Exporting MoGe-2 {debug_format.upper()}") + export_moge_debug_mesh( + output_directory, + points[0], + depths[0], + masks[0], + normals[0], + source_images[0], + debug_format, + ) remove_background_model_dir = getattr( args, "remove_background_model_dir", "" @@ -1614,7 +1717,14 @@ def run_moge(args, model_cache=None, cancel_check=None): model_cache=model_cache, cancel_check=cancel_check, ) - foreground.save(output_directory / "foreground.png") + foreground_path = output_directory / "foreground.png" + foreground.save(foreground_path) + if points: + foreground_mask = foreground_mesh_mask( + foreground_path, + result["depth"][0].shape, + ) + result["mask"][0] &= foreground_mask else: infer_foreground_frames( images, @@ -1628,6 +1738,19 @@ def run_moge(args, model_cache=None, cancel_check=None): model_cache=model_cache, cancel_check=cancel_check, ) + if bool(getattr(args, "surface_mesh", False)): + if len(images) != 1 or not points: + raise RuntimeError("MoGe-2 Surface requires one point map") + write_status(args.status, 0.94, "Building official MoGe-2 surface") + export_moge_surface_archive( + output_directory / "mesh.npz", + points[0], + result["depth"][0], + result["mask"][0], + result["normal"][0], + result["intrinsics"][0], + ) + np.savez(output_directory / "results.npz", **result) copy_frames(images, output_directory) write_status( args.status, diff --git a/tests/test_addon.py b/tests/test_addon.py index 0d0575b..ea54b63 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -87,8 +87,6 @@ class AddonTest(unittest.TestCase): self.context_menu_draws.clear() self.outliner_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), ) @@ -732,11 +730,8 @@ class AddonTest(unittest.TestCase): ], ) - def test_conversion_dialog_remembers_last_depth_model(self): + def test_conversion_dialog_leaves_last_properties_to_blender(self): operator = self.addon.ConvertImageToDepthPlane() - operator.properties = SimpleNamespace( - is_property_set=lambda _name: False, - ) context = SimpleNamespace( object=SimpleNamespace( data=SimpleNamespace(source="FILE", frame_duration=1) @@ -749,16 +744,9 @@ class AddonTest(unittest.TestCase): ), ) - self.addon.image_depth._last_depth_plane_model = "GIANT_1_1" - with patch.object( - self.addon.client, - "model_ready", - side_effect=lambda key: key == "GIANT_1_1", - ): - result = operator.invoke(context, None) + result = operator.invoke(context, None) self.assertIs(result, operator) - self.assertNotEqual(getattr(operator, "snapform_model", ""), "GIANT_1_1") def test_remove_background_defaults_to_empty_image_frames(self): source_object = SimpleNamespace( @@ -995,10 +983,6 @@ class AddonTest(unittest.TestCase): self.assertTrue(calls[0]["snapform_unshaded"]) self.assertEqual(calls[0]["snapform_depth_data_name"], "img.depth") self.assertTrue(calls[0]["snapform_delete_input"]) - self.assertEqual( - self.addon.image_depth._last_depth_plane_model, - "GIANT_1_1", - ) def test_packed_image_export_writes_a_png_copy(self): class PackedImage: diff --git a/tools/build_geometry_nodes_asset.py b/tools/build_geometry_nodes_asset.py index 466a3ec..b217f89 100644 --- a/tools/build_geometry_nodes_asset.py +++ b/tools/build_geometry_nodes_asset.py @@ -4,6 +4,7 @@ import bpy GROUP_NAME = "O DA3 Surface" +MOGE_SURFACE_GROUP_NAME = "O MoGe2 Surface" DEPTH_PLANE_GROUP_NAME = "O DA3 Depth Plane" MOGE_DEPTH_PLANE_GROUP_NAME = "O MoGe2 Depth Plane" MATERIAL_DISPLACEMENT_ATTRIBUTE = "o_material_displacement" @@ -357,8 +358,13 @@ def add_uv_reprojection(group, source, geometry, location): return store_uv.outputs["Geometry"] -def build_group(): - group = bpy.data.node_groups.new(GROUP_NAME, "GeometryNodeTree") +def build_group( + group_name=GROUP_NAME, + include_confidence=True, + include_stretch=True, + include_mask=False, +): + group = bpy.data.node_groups.new(group_name, "GeometryNodeTree") group.is_modifier = True interface = group.interface geometry_input_socket = new_socket( @@ -393,37 +399,50 @@ def build_group(): "INPUT", "NodeSocketFloat", ) - smooth_amount.default_value = 1.0 + smooth_amount.default_value = 0.5 smooth_amount.min_value = 0.0 smooth_amount.max_value = 1.0 smooth_amount.subtype = "FACTOR" + if include_mask: + mask_socket = new_socket( + interface, + "Mask", + "INPUT", + "NodeSocketBool", + ) + mask_socket.default_value = True + mask_socket.structure_type = "SINGLE" - 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 = 100.0 - confidence.subtype = "PERCENTAGE" - edge = new_socket( - interface, - "Stretch", - "INPUT", - "NodeSocketFloat", - parent=culling_panel, - ) - edge.default_value = 10.0 - edge.min_value = 0.0 - edge.max_value = 100.0 - edge.subtype = "PERCENTAGE" + culling_panel = None + if include_confidence or include_stretch: + culling_panel = interface.new_panel( + name="Culling", + default_closed=False, + ) + if include_confidence: + confidence = new_socket( + interface, + "Confidence", + "INPUT", + "NodeSocketFloat", + parent=culling_panel, + ) + confidence.default_value = 0.0 + confidence.min_value = 0.0 + confidence.max_value = 100.0 + confidence.subtype = "PERCENTAGE" + if include_stretch: + edge = new_socket( + interface, + "Stretch", + "INPUT", + "NodeSocketFloat", + parent=culling_panel, + ) + edge.default_value = 10.0 + edge.min_value = 0.0 + edge.max_value = 100.0 + edge.subtype = "PERCENTAGE" options_panel = interface.new_panel( name="Options", @@ -444,8 +463,16 @@ def build_group(): 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"}) + confidence_input = None + if include_confidence: + confidence_input = new_group_input( + nodes, + (-1800.0, 800.0), + {"Confidence"}, + ) + stretch_input = None + if include_stretch: + stretch_input = new_group_input(nodes, (-1800.0, 50.0), {"Stretch"}) remesh_input = new_group_input( nodes, (-100.0, -500.0), @@ -457,89 +484,71 @@ def build_group(): 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_scale = nodes.new("ShaderNodeMath") - confidence_scale.operation = "MULTIPLY" - confidence_scale.inputs[1].default_value = 0.01 - confidence_scale.location = (-1600.0, 800.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"], - confidence_scale.inputs[0], - ) - links.new( - confidence_scale.outputs[0], - input_socket(confidence_compare, "B", 1), - ) + geometry = geometry_input.outputs["Geometry"] + if include_confidence: + 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_scale = nodes.new("ShaderNodeMath") + confidence_scale.operation = "MULTIPLY" + confidence_scale.inputs[1].default_value = 0.01 + confidence_scale.location = (-1600.0, 800.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"], + confidence_scale.inputs[0], + ) + links.new( + confidence_scale.outputs[0], + input_socket(confidence_compare, "B", 1), + ) + delete_confidence = nodes.new("GeometryNodeDeleteGeometry") + delete_confidence.domain = "FACE" + delete_confidence.location = (-1050.0, 600.0) + links.new(geometry, delete_confidence.inputs["Geometry"]) + links.new( + confidence_compare.outputs["Result"], + delete_confidence.inputs["Selection"], + ) + geometry = delete_confidence.outputs["Geometry"] - 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_scale = nodes.new("ShaderNodeMath") - edge_scale.operation = "MULTIPLY" - edge_scale.inputs[1].default_value = 0.01 - edge_scale.location = (-1600.0, -50.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_scale.inputs[0]) - links.new(edge_scale.outputs[0], 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)) - - 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]) - full_stretch = new_compare( - nodes, - "FLOAT", - "GREATER_EQUAL", - (-950.0, -20.0), - 0.999999, - ) - links.new( - edge_scale.outputs[0], - input_socket(full_stretch, "A", 0), - ) - culling_selection = nodes.new("FunctionNodeBooleanMath") - culling_selection.operation = "OR" - culling_selection.location = (-550.0, 300.0) - links.new(first_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"]) + culled_geometry = geometry + if include_stretch: + 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_scale = nodes.new("ShaderNodeMath") + edge_scale.operation = "MULTIPLY" + edge_scale.inputs[1].default_value = 0.01 + edge_scale.location = (-1600.0, -50.0) + edge_compare = new_compare( + nodes, + "FLOAT", + "GREATER_THAN", + (-1200.0, 350.0), + 0.0, + ) + links.new(stretch_input.outputs["Stretch"], edge_scale.inputs[0]) + links.new(edge_attribute.outputs["Attribute"], input_socket(edge_compare, "A", 0)) + links.new(edge_scale.outputs[0], input_socket(edge_compare, "B", 1)) + delete_geometry = nodes.new("GeometryNodeDeleteGeometry") + delete_geometry.domain = "EDGE" + delete_geometry.location = (-350.0, 100.0) + links.new(geometry, delete_geometry.inputs["Geometry"]) + links.new(edge_compare.outputs["Result"], delete_geometry.inputs["Selection"]) + culled_geometry = delete_geometry.outputs["Geometry"] vertex_neighbors = nodes.new("GeometryNodeInputMeshVertexNeighbors") vertex_neighbors.location = (-550.0, -320.0) @@ -557,7 +566,7 @@ def build_group(): 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(culled_geometry, delete_orphans.inputs["Geometry"]) links.new(orphan_vertex.outputs["Result"], delete_orphans.inputs["Selection"]) remeshed = add_remesh_zone( @@ -622,6 +631,49 @@ def build_group(): return group +def configure_moge_surface_group(group): + nodes = group.nodes + links = group.links + mask_input = new_group_input(nodes, (-1800.0, 1050.0), {"Mask"}) + mask_attribute = nodes.new("GeometryNodeInputNamedAttribute") + mask_attribute.data_type = "BOOLEAN" + mask_attribute.inputs["Name"].default_value = "o_mask" + mask_attribute.location = (-1550.0, 1050.0) + invert_mask = nodes.new("FunctionNodeBooleanMath") + invert_mask.operation = "NOT" + invert_mask.location = (-1350.0, 1050.0) + use_mask = nodes.new("FunctionNodeBooleanMath") + use_mask.operation = "AND" + use_mask.location = (-1150.0, 1050.0) + links.new(mask_attribute.outputs["Attribute"], invert_mask.inputs[0]) + links.new(mask_input.outputs["Mask"], use_mask.inputs[0]) + links.new(invert_mask.outputs[0], use_mask.inputs[1]) + first_processing = min( + ( + node + for node in nodes + if node.bl_idname == "GeometryNodeDeleteGeometry" + ), + key=lambda node: node.location.x, + ) + geometry_input = first_processing.inputs["Geometry"] + previous_link = geometry_input.links[0] + previous_geometry = previous_link.from_socket + links.remove(previous_link) + delete_mask = nodes.new("GeometryNodeDeleteGeometry") + delete_mask.domain = "FACE" + delete_mask.location = (-750.0, 650.0) + links.new(previous_geometry, delete_mask.inputs["Geometry"]) + links.new(use_mask.outputs[0], delete_mask.inputs["Selection"]) + links.new(delete_mask.outputs["Geometry"], geometry_input) + group.description = ( + "Cull MoGe-2 stretched faces, repair topology, preserve boundaries, " + "and reproject UVs." + ) + group.use_fake_user = True + return group + + def build_sequence_frame(nodes, links, scene_frame, frames_input, y): relative_frame = nodes.new("ShaderNodeMath") relative_frame.operation = "SUBTRACT" @@ -1246,6 +1298,13 @@ def main(): bpy.ops.wm.read_factory_settings(use_empty=True) bpy.context.preferences.filepaths.save_version = 0 build_group() + moge_surface_group = build_group( + MOGE_SURFACE_GROUP_NAME, + include_confidence=False, + include_stretch=False, + include_mask=True, + ) + configure_moge_surface_group(moge_surface_group) build_depth_plane_group() build_moge_depth_plane_group() OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True) diff --git a/tools/validate_depth_plane_asset.py b/tools/validate_depth_plane_asset.py index 2fb44c8..8ac4544 100644 --- a/tools/validate_depth_plane_asset.py +++ b/tools/validate_depth_plane_asset.py @@ -20,6 +20,13 @@ from snapform.image_mesh import load_geometry_node_group def main(): bpy.ops.wm.read_factory_settings(use_empty=True) group = load_depth_plane_node_group() + geometry_textures = [ + node + for node in group.nodes + if node.bl_idname == "GeometryNodeImageTexture" + ] + assert geometry_textures + assert all(node.extension == "EXTEND" for node in geometry_textures) background_depth = next( item for item in group.interface.items_tree @@ -212,6 +219,11 @@ def main(): node.bl_idname in forbidden_outer_nodes for node in material.node_tree.nodes ) + assert all( + node.extension == "EXTEND" + for node in material.node_tree.nodes + if node.bl_idname == "ShaderNodeTexImage" + ) surface_material = create_material("Surface", foreground, 1.0) surface_group = bpy.data.node_groups[SURFACE_MATERIAL_NODE_GROUP_NAME] surface_group_nodes = [ @@ -251,6 +263,11 @@ def main(): node.bl_idname in forbidden_outer_nodes for node in surface_material.node_tree.nodes ) + assert all( + node.extension == "EXTEND" + for node in surface_material.node_tree.nodes + if node.bl_idname == "ShaderNodeTexImage" + ) modifier = add_depth_plane_processing( obj, depth, diff --git a/tools/validate_geometry_nodes_asset.py b/tools/validate_geometry_nodes_asset.py index 6f3964c..6d69bed 100644 --- a/tools/validate_geometry_nodes_asset.py +++ b/tools/validate_geometry_nodes_asset.py @@ -22,9 +22,8 @@ def set_input(modifier, group, name, 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) + confidence = mesh.attributes.new("o_face_confidence", "FLOAT", "FACE") + confidence.data.foreach_set("value", [1.0] * count) background = mesh.attributes.new("o_background_face", "BOOLEAN", "FACE") background.data.foreach_set( "value", @@ -56,7 +55,7 @@ def main(): visible_outputs = [ output for output in node.outputs if output.name and not output.hide ] - assert 1 <= len(visible_outputs) <= 3 + assert 1 <= len(visible_outputs) <= 4 uv_stores = [ node for node in group.nodes @@ -79,6 +78,14 @@ def main(): assert [item.name for item in panels["Options"].interface_items] == [ "Fix Topology", ] + smooth = next( + item + for item in group.interface.items_tree + if item.item_type == "SOCKET" + and item.in_out == "INPUT" + and item.name == "Smooth" + ) + assert smooth.default_value == 0.5 for name, default_value in (("Confidence", 0.0), ("Stretch", 10.0)): item = next( item @@ -98,6 +105,8 @@ def main(): [], [(0, 2, 1), (1, 2, 3)], ) + edge_stretch = mesh.attributes.new("o_edge_stretch", "FLOAT", "EDGE") + edge_stretch.data.foreach_set("value", [0.05] * len(mesh.edges)) 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]) @@ -121,13 +130,13 @@ def main(): finally: evaluated.to_mesh_clear() - set_input(modifier, group, "Stretch", 100.0) + set_input(modifier, group, "Stretch", 0.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] == 100.0 + assert modifier[stretch_item.identifier] == 0.0 obj.update_tag(refresh={"DATA"}) bpy.context.view_layer.update() culled = obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) diff --git a/tools/validate_moge_depth_plane.py b/tools/validate_moge_depth_plane.py index 240b9b8..8f7fb74 100644 --- a/tools/validate_moge_depth_plane.py +++ b/tools/validate_moge_depth_plane.py @@ -72,6 +72,21 @@ def main(): material_group = bpy.data.node_groups[MOGE_MATERIAL_NODE_GROUP_NAME] assert not any(node.bl_idname == "ShaderNodeNormalMap" for node in material_group.nodes) assert any(node.bl_idname == "ShaderNodeBump" for node in material_group.nodes) + material = plane.data.materials[0] + image_textures = [ + node + for node in material.node_tree.nodes + if node.bl_idname == "ShaderNodeTexImage" + ] + assert image_textures + assert all(node.extension == "EXTEND" for node in image_textures) + geometry_textures = [ + node + for node in geometry_group.nodes + if node.bl_idname == "GeometryNodeImageTexture" + ] + assert geometry_textures + assert all(node.extension == "EXTEND" for node in geometry_textures) print("MoGe-2 depth plane validation passed") diff --git a/tools/validate_moge_surface.py b/tools/validate_moge_surface.py new file mode 100644 index 0000000..fe76906 --- /dev/null +++ b/tools/validate_moge_surface.py @@ -0,0 +1,176 @@ +import sys +import tempfile +from pathlib import Path + +import bpy +import numpy as np + + +PROJECT_ROOT = Path(__file__).parents[1] +sys.path.insert(0, str(PROJECT_ROOT / "src")) + +from snapform.image_depth import MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME +from snapform.image_mesh import ( + MOGE_NODE_GROUP_NAME, + build_image_mesh, + prepare_moge_surface_archive, +) + + +def main(): + bpy.ops.wm.read_factory_settings(use_empty=True) + arguments = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else [] + if arguments: + validate_real_result(Path(arguments[0]), Path(arguments[1])) + return + height, width = 48, 64 + u, v = np.meshgrid( + np.linspace(-0.5, 0.5, width, dtype=np.float32), + np.linspace(-0.375, 0.375, height, dtype=np.float32), + ) + depth = np.full((1, height, width), 2.0, dtype=np.float32) + points = np.stack((u * 2.0, v * 2.0, depth[0]), axis=-1)[None] + mask = np.ones((1, height, width), dtype=bool) + mask[:, :2] = False + intrinsics = np.array( + [[[1.0, 0.0, 0.5], [0.0, 1.0, 0.5], [0.0, 0.0, 1.0]]], + dtype=np.float32, + ) + image = bpy.data.images.new("MoGe Surface Test", width=width, height=height) + source = bpy.data.objects.new("MoGe Surface Test", None) + source.empty_display_type = "IMAGE" + source.data = image + bpy.context.collection.objects.link(source) + with tempfile.TemporaryDirectory() as temporary: + directory = Path(temporary) + results_path = directory / "results.npz" + mesh_path = directory / "mesh.npz" + np.savez( + results_path, + depth=depth, + mask=mask, + points=points, + intrinsics=intrinsics, + ) + prepare_moge_surface_archive(results_path, mesh_path) + surface = build_image_mesh( + bpy.context, + source, + mesh_path, + False, + True, + False, + False, + True, + ) + assert len(surface.data.vertices) > 0 + assert surface.modifiers[0].node_group.name == MOGE_NODE_GROUP_NAME + inputs = { + item.name + for item in surface.modifiers[0].node_group.interface.items_tree + if getattr(item, "in_out", None) == "INPUT" + } + assert "Stretch" not in inputs + assert "Mask" in inputs + assert "Confidence" not in inputs + assert "o_mask" in surface.data.attributes + assert surface.data.attributes["o_mask"].domain == "FACE" + attribute_names = set(surface.data.attributes.keys()) + assert attribute_names >= {"o_depth", "o_mask"} + assert "o_edge_stretch" not in attribute_names + assert not { + "o_confidence", + "o_confidence_percentile", + "o_face_confidence", + "o_background", + "o_background_face", + }.intersection(attribute_names) + mask_socket = next( + item + for item in surface.modifiers[0].node_group.interface.items_tree + if getattr(item, "in_out", None) == "INPUT" and item.name == "Mask" + ) + assert mask_socket.default_value is True + assert mask_socket.structure_type == "SINGLE" + assert not mask_socket.parent.name + assert not any( + item.item_type == "PANEL" and item.name == "Culling" + for item in surface.modifiers[0].node_group.interface.items_tree + ) + smooth_socket = next( + item + for item in surface.modifiers[0].node_group.interface.items_tree + if getattr(item, "in_out", None) == "INPUT" and item.name == "Smooth" + ) + assert smooth_socket.default_value == 0.5 + assert not any( + node.bl_idname == "GeometryNodeInputNamedAttribute" + and node.inputs["Name"].default_value == "o_face_confidence" + for node in surface.modifiers[0].node_group.nodes + ) + assert MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME in bpy.data.node_groups + material_group = bpy.data.node_groups[MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME] + assert any(node.bl_idname == "ShaderNodeBump" for node in material_group.nodes) + assert any(child.type == "CAMERA" for child in surface.children) + print("MoGe-2 surface validation passed") + + +def validate_real_result(results_path, image_path): + image = bpy.data.images.load(str(image_path), check_existing=False) + source = bpy.data.objects.new("MoGe Real Surface", None) + source.empty_display_type = "IMAGE" + source.data = image + bpy.context.collection.objects.link(source) + mesh_path = results_path.with_name("mesh.npz") + if not mesh_path.is_file(): + prepare_moge_surface_archive(results_path, mesh_path) + surface = build_image_mesh( + bpy.context, + source, + mesh_path, + False, + True, + False, + False, + True, + ) + assert len(surface.data.vertices) > 100 + assert len(surface.data.polygons) > 100 + assert surface.modifiers[0].node_group.name == MOGE_NODE_GROUP_NAME + assert MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME in bpy.data.node_groups + assert surface.data.has_custom_normals + normals = np.empty(len(surface.data.corner_normals) * 3, dtype=np.float32) + surface.data.corner_normals.foreach_get("vector", normals) + normals = normals.reshape(-1, 3) + assert np.isfinite(normals).all() + assert np.min(np.linalg.norm(normals, axis=1)) > 0.99 + material_group = bpy.data.node_groups[MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME] + assert any(node.bl_idname == "ShaderNodeBump" for node in material_group.nodes) + attribute_names = set(surface.data.attributes.keys()) + assert {"o_depth", "o_mask"}.issubset(attribute_names) + assert "o_edge_stretch" not in attribute_names + assert not { + "o_confidence", + "o_confidence_percentile", + "o_face_confidence", + "o_background", + "o_background_face", + }.intersection(attribute_names) + camera = next(child for child in surface.children if child.type == "CAMERA") + dimensions = surface.dimensions + assert max(dimensions) < 10.0 + assert min(dimensions.x, dimensions.y) > 0.1 + assert 1.0 <= camera.data.lens <= 500.0 + assert abs(camera.data.shift_x) <= 1.0 + assert abs(camera.data.shift_y) <= 1.0 + print( + f"Real MoGe-2 surface validation passed: " + f"{len(surface.data.vertices)} vertices, {len(surface.data.polygons)} faces, " + f"dimensions={tuple(round(value, 4) for value in dimensions)}, " + f"lens={camera.data.lens:.3f}, shift=({camera.data.shift_x:.4f}, " + f"{camera.data.shift_y:.4f})" + ) + + +if __name__ == "__main__": + main()