feat: Add surface generation and media handling
- 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.
This commit is contained in:
@@ -8,13 +8,37 @@ 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,
|
||||
prepare_moge_surface_archive,
|
||||
)
|
||||
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():
|
||||
@@ -37,26 +61,38 @@ def main():
|
||||
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)
|
||||
results_path = directory / "results.npz"
|
||||
mesh_path = directory / "mesh.npz"
|
||||
np.savez(
|
||||
results_path,
|
||||
depth=depth,
|
||||
mask=mask,
|
||||
points=points,
|
||||
intrinsics=intrinsics,
|
||||
normal=np.broadcast_to(
|
||||
np.array([0.0, 0.0, -1.0], dtype=np.float32),
|
||||
(1, height, width, 3),
|
||||
),
|
||||
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",
|
||||
)
|
||||
prepare_moge_surface_archive(results_path, mesh_path)
|
||||
surface = build_image_mesh(
|
||||
bpy.context,
|
||||
source,
|
||||
@@ -65,8 +101,11 @@ def main():
|
||||
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
|
||||
@@ -88,7 +127,7 @@ def main():
|
||||
assert not {
|
||||
"o_confidence",
|
||||
"o_confidence_percentile",
|
||||
"o_face_confidence",
|
||||
"o_confidence",
|
||||
"o_background",
|
||||
"o_background_face",
|
||||
}.intersection(attribute_names)
|
||||
@@ -134,7 +173,7 @@ def main():
|
||||
assert smooth_socket.default_value == 0.5
|
||||
assert not any(
|
||||
node.bl_idname == "GeometryNodeInputNamedAttribute"
|
||||
and node.inputs["Name"].default_value == "o_face_confidence"
|
||||
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
|
||||
@@ -154,6 +193,11 @@ def main():
|
||||
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"
|
||||
@@ -172,7 +216,10 @@ def validate_real_result(results_path, image_path):
|
||||
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)
|
||||
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,
|
||||
@@ -181,6 +228,7 @@ def validate_real_result(results_path, image_path):
|
||||
True,
|
||||
False,
|
||||
True,
|
||||
normal_path=normal_path,
|
||||
)
|
||||
assert len(surface.data.vertices) > 100
|
||||
assert len(surface.data.polygons) > 100
|
||||
@@ -202,7 +250,7 @@ def validate_real_result(results_path, image_path):
|
||||
assert not {
|
||||
"o_confidence",
|
||||
"o_confidence_percentile",
|
||||
"o_face_confidence",
|
||||
"o_confidence",
|
||||
"o_background",
|
||||
"o_background_face",
|
||||
}.intersection(attribute_names)
|
||||
|
||||
Reference in New Issue
Block a user