- Implemented surface generation functionality in `surface.py`, including geometry estimation and texture writing. - Introduced media input handling in `input.py`, supporting image and video frame extraction. - Added model management for geometry adapters in `models/__init__.py`. - Created BEN2 and DA3 model inference logic in `ben2.py` and `da3.py`, respectively. - Developed ONNX model inference utilities for BEN2 and DA3 in `onnx_ben2.py` and `onnx_da3.py`. - Added support for MoGe-2 model inference in `moge2.py` with necessary ONNX utilities. - Implemented output handling for depth and surface data in `depth.py` and `surface.py`. - Created debug output utilities for DA3 and MoGe-2 predictions in `debug_da3.py` and `debug_moge2.py`. - Enhanced ONNX runtime session management in `onnx_runtime.py`. - Updated `__init__.py` files to include new modules and functions for better organization.
100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
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"))
|
|
sys.path.insert(
|
|
0,
|
|
str(next((PROJECT_ROOT / "src" / "snapform" / "wheels").glob("blendjob-*.whl"))),
|
|
)
|
|
|
|
from snapform.blender.depth_plane import (
|
|
MOGE_DEPTH_PLANE_NODE_GROUP_NAME,
|
|
MOGE_MATERIAL_NODE_GROUP_NAME,
|
|
build_moge2_depth_plane,
|
|
)
|
|
|
|
|
|
def main():
|
|
bpy.ops.wm.read_factory_settings(use_empty=True)
|
|
image = bpy.data.images.new("MoGe Test", width=64, height=48, alpha=True)
|
|
source = bpy.data.objects.new("MoGe Test", None)
|
|
source.empty_display_type = "IMAGE"
|
|
source.data = image
|
|
bpy.context.collection.objects.link(source)
|
|
depth = np.linspace(1.0, 2.0, 48 * 64, dtype=np.float32).reshape(1, 48, 64)
|
|
normal = np.zeros((1, 48, 64, 3), dtype=np.float32)
|
|
normal[..., 2] = -1.0
|
|
mask = np.ones((1, 48, 64), dtype=bool)
|
|
mask[:, :5] = 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,
|
|
)
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
path = Path(temporary) / "results.npz"
|
|
np.savez(
|
|
path,
|
|
depth=depth,
|
|
normal=normal,
|
|
mask=mask,
|
|
intrinsics=intrinsics,
|
|
)
|
|
plane = build_moge2_depth_plane(
|
|
bpy.context,
|
|
source,
|
|
path,
|
|
delete_source=True,
|
|
)
|
|
assert plane.modifiers[0].node_group.name == MOGE_DEPTH_PLANE_NODE_GROUP_NAME
|
|
assert MOGE_MATERIAL_NODE_GROUP_NAME in bpy.data.node_groups
|
|
geometry_group = plane.modifiers[0].node_group
|
|
inputs = {
|
|
item.name
|
|
for item in geometry_group.interface.items_tree
|
|
if getattr(item, "in_out", None) == "INPUT"
|
|
}
|
|
assert "Displacement" in inputs
|
|
assert "Background Depth" in inputs
|
|
assert "Mask Image" not in inputs
|
|
assert "Use Mask" not in inputs
|
|
sockets = {
|
|
item.name: item
|
|
for item in geometry_group.interface.items_tree
|
|
if getattr(item, "in_out", None) == "INPUT"
|
|
}
|
|
assert sockets["Subdivision"].default_value == 6
|
|
assert sockets["Depth Scale"].structure_type == "SINGLE"
|
|
assert sockets["Background Depth"].structure_type == "SINGLE"
|
|
assert sockets["Background Depth"].subtype == "DISTANCE"
|
|
assert sockets["Background Depth"].default_value == 1.0
|
|
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)
|
|
assert any(node.image.name == "MoGe Test.depth" 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")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|