nan
835225d264
- 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.
275 lines
10 KiB
Python
275 lines
10 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_SURFACE_MATERIAL_NODE_GROUP_NAME
|
|
from snapform.blender.surface import (
|
|
MOGE_NODE_GROUP_NAME,
|
|
build_image_mesh,
|
|
)
|
|
from snapform.server.geometry import GeometryFrame, build_surface
|
|
from snapform.server.outputs.surface import write_surface_archive
|
|
|
|
|
|
def write_test_normal_texture(normal, output_path):
|
|
converted = normal * np.array([1.0, -1.0, -1.0], dtype=np.float32)
|
|
rgb = np.clip((converted + 1.0) * 0.5, 0.0, 1.0)
|
|
alpha = np.ones((*rgb.shape[:2], 1), dtype=np.float32)
|
|
pixels = np.concatenate((rgb, alpha), axis=-1)
|
|
image = bpy.data.images.new(
|
|
"MoGe Test Normal",
|
|
width=rgb.shape[1],
|
|
height=rgb.shape[0],
|
|
alpha=True,
|
|
float_buffer=False,
|
|
)
|
|
image.pixels.foreach_set(pixels.reshape(-1))
|
|
image.file_format = "PNG"
|
|
image.filepath_raw = str(output_path)
|
|
image.save()
|
|
return output_path
|
|
|
|
|
|
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)
|
|
foreground = bpy.data.images.new(
|
|
"MoGe Surface Test.foreground",
|
|
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)
|
|
mesh_path = directory / "mesh.npz"
|
|
normal = np.broadcast_to(
|
|
np.array([0.0, 0.0, -1.0], dtype=np.float32),
|
|
(height, width, 3),
|
|
)
|
|
pixel_intrinsics = intrinsics[0].copy()
|
|
pixel_intrinsics[0] *= width
|
|
pixel_intrinsics[1] *= height
|
|
pixel_intrinsics[2] = (0.0, 0.0, 1.0)
|
|
frame = GeometryFrame(
|
|
depth=depth[0],
|
|
points=points[0],
|
|
intrinsics=pixel_intrinsics,
|
|
valid_mask=np.isfinite(depth[0]) & (depth[0] > 0.0),
|
|
model_mask=mask[0],
|
|
)
|
|
write_surface_archive(build_surface(frame), mesh_path)
|
|
normal_path = write_test_normal_texture(
|
|
normal,
|
|
directory / "normal.png",
|
|
)
|
|
surface = build_image_mesh(
|
|
bpy.context,
|
|
source,
|
|
mesh_path,
|
|
False,
|
|
True,
|
|
False,
|
|
True,
|
|
color_image=foreground,
|
|
normal_path=normal_path,
|
|
)
|
|
assert len(surface.data.vertices) > 0
|
|
assert all(len(polygon.vertices) == 3 for polygon in surface.data.polygons)
|
|
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 "Depth Continuity" 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"
|
|
assert not all(item.value for item in surface.data.attributes["o_mask"].data)
|
|
attribute_names = set(surface.data.attributes.keys())
|
|
assert attribute_names >= {"o_depth", "o_mask"}
|
|
assert "o_depth_discontinuity" in attribute_names
|
|
assert surface.data.attributes["o_depth_discontinuity"].data_type == "FLOAT"
|
|
assert "o_edge_stretch" not in attribute_names
|
|
assert not {
|
|
"o_confidence",
|
|
"o_confidence_percentile",
|
|
"o_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 False
|
|
assert mask_socket.structure_type == "SINGLE"
|
|
assert mask_socket.parent.name == "Culling"
|
|
depth_edge_socket = next(
|
|
item
|
|
for item in surface.modifiers[0].node_group.interface.items_tree
|
|
if getattr(item, "in_out", None) == "INPUT"
|
|
and item.name == "Depth Continuity"
|
|
)
|
|
assert depth_edge_socket.default_value == 50.0
|
|
assert depth_edge_socket.subtype == "PERCENTAGE"
|
|
assert depth_edge_socket.structure_type == "SINGLE"
|
|
culling_panel = next(
|
|
item for item in surface.modifiers[0].node_group.interface.items_tree
|
|
if item.item_type == "PANEL" and item.name == "Culling"
|
|
)
|
|
assert [item.name for item in culling_panel.interface_items] == [
|
|
"Depth Continuity",
|
|
"Mask",
|
|
]
|
|
input_sockets = [
|
|
item for item in surface.modifiers[0].node_group.interface.items_tree
|
|
if item.item_type == "SOCKET" and item.in_out == "INPUT"
|
|
]
|
|
assert all(item.structure_type == "SINGLE" for item in input_sockets)
|
|
resolution_socket = next(
|
|
item for item in input_sockets if item.name == "Resolution"
|
|
)
|
|
assert abs(resolution_socket.default_value - 0.1) < 1e-6
|
|
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_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)
|
|
normal_map = next(
|
|
node for node in material_group.nodes
|
|
if node.bl_idname == "ShaderNodeNormalMap"
|
|
)
|
|
assert normal_map.space == "OBJECT"
|
|
normal_strength = next(
|
|
item for item in material_group.interface.items_tree
|
|
if getattr(item, "in_out", None) == "INPUT"
|
|
and item.name == "Normal Strength"
|
|
)
|
|
assert normal_strength.default_value == 1.0
|
|
assert normal_strength.min_value == 0.0
|
|
assert normal_map.inputs["Strength"].is_linked
|
|
assert not surface.data.has_custom_normals
|
|
color_texture = next(
|
|
node for node in surface.data.materials[0].node_tree.nodes
|
|
if node.bl_idname == "ShaderNodeTexImage" and node.image == foreground
|
|
)
|
|
assert color_texture.image.name == "MoGe Surface Test.foreground"
|
|
normal_texture = next(
|
|
node for node in surface.data.materials[0].node_tree.nodes
|
|
if node.bl_idname == "ShaderNodeTexImage"
|
|
and node.image.colorspace_settings.name == "Non-Color"
|
|
)
|
|
assert normal_texture.image.name == "MoGe Surface Test.normal"
|
|
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():
|
|
raise ValueError(f"Surface archive is missing: {mesh_path}")
|
|
normal_path = results_path.with_name("normal.png")
|
|
if not normal_path.is_file():
|
|
raise ValueError(f"Surface normal texture is missing: {normal_path}")
|
|
surface = build_image_mesh(
|
|
bpy.context,
|
|
source,
|
|
mesh_path,
|
|
False,
|
|
True,
|
|
False,
|
|
True,
|
|
normal_path=normal_path,
|
|
)
|
|
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 not surface.data.has_custom_normals
|
|
material_group = bpy.data.node_groups[MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME]
|
|
assert any(node.bl_idname == "ShaderNodeBump" for node in material_group.nodes)
|
|
normal_map = next(
|
|
node for node in material_group.nodes
|
|
if node.bl_idname == "ShaderNodeNormalMap"
|
|
)
|
|
assert normal_map.space == "OBJECT"
|
|
assert normal_map.inputs["Strength"].is_linked
|
|
attribute_names = set(surface.data.attributes.keys())
|
|
assert {"o_depth", "o_mask"}.issubset(attribute_names)
|
|
assert "o_depth_discontinuity" in attribute_names
|
|
assert "o_edge_stretch" not in attribute_names
|
|
assert not {
|
|
"o_confidence",
|
|
"o_confidence_percentile",
|
|
"o_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()
|