icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Add surface generation and media handling

Browse Source 
- 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:
马南
2026-08-18 23:27:10 +08:00
parent 9dfa49602c
commit 835225d264
62 changed files with 2126 additions and 2318 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
tools/build_geometry_nodes_asset.py
View File

@@ -381,6 +381,15 @@ def build_group(
)
geometry_input_socket.structure_type = "SINGLE"
remesh_iterations = new_socket(
interface,
"Iterations",
"INPUT",
"NodeSocketInt",
)
remesh_iterations.default_value = 2
remesh_iterations.min_value = 0
remesh_iterations.max_value = 8
resolution = new_socket(
interface,
"Resolution",
@@ -390,15 +399,6 @@ def build_group(
resolution.default_value = 0.1
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",
@@ -485,7 +485,7 @@ def build_group(
remesh_input = new_group_input(
nodes,
(-100.0, -500.0),
{"Resolution", "Relax Iterations", "Smooth"},
{"Resolution", "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"})
@@ -497,7 +497,7 @@ def build_group(
if include_confidence:
confidence_attribute = nodes.new("GeometryNodeInputNamedAttribute")
confidence_attribute.data_type = "FLOAT"
confidence_attribute.inputs["Name"].default_value = "o_face_confidence"
confidence_attribute.inputs["Name"].default_value = "o_confidence"
confidence_attribute.location = (-1800.0, 600.0)
confidence_scale = nodes.new("ShaderNodeMath")
confidence_scale.operation = "MULTIPLY"
@@ -595,7 +595,7 @@ def build_group(
group,
culled_geometry,
remesh_input.outputs["Resolution"],
remesh_input.outputs["Relax Iterations"],
remesh_input.outputs["Iterations"],
topology_input.outputs["Fix Topology"],
remesh_input.outputs["Smooth"],
(-300.0, 100.0),

7
tools/validate_depth_plane_asset.py
View File

@@ -4,7 +4,12 @@ import sys
import bpy
sys.path.insert(0, str(Path(__file__).parents[1] / "src"))
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 (
DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME,

23
tools/validate_geometry_nodes_asset.py
View File

@@ -5,7 +5,12 @@ import bpy
import numpy as np
sys.path.insert(0, str(Path(__file__).parents[1] / "src"))
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 (
get_modifier_input,
@@ -26,7 +31,7 @@ def set_input(modifier, group, name, value):
def add_face_attributes(mesh, count):
confidence = mesh.attributes.new("o_face_confidence", "FLOAT", "FACE")
confidence = mesh.attributes.new("o_confidence", "FLOAT", "FACE")
confidence.data.foreach_set("value", [1.0] * count)
@@ -48,6 +53,15 @@ def main():
if item.item_type == "SOCKET" and item.in_out == "INPUT"
]
assert all(item.structure_type == "SINGLE" for item in input_sockets)
assert [item.name for item in input_sockets[:4]] == [
"Geometry",
"Iterations",
"Resolution",
"Smooth",
]
iterations = next(item for item in input_sockets if item.name == "Iterations")
assert iterations.default_value == 2
assert iterations.min_value == 0
resolution = next(item for item in input_sockets if item.name == "Resolution")
assert abs(resolution.default_value - 0.1) < 1e-6
assert group.asset_data is None
@@ -123,10 +137,11 @@ def main():
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)
assert "o_confidence" in mesh.attributes
assert not {
"o_background",
"o_background_face",
"o_confidence",
"o_face_confidence",
"o_confidence_percentile",
}.intersection(mesh.attributes.keys())
@@ -137,7 +152,7 @@ def main():
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, "Iterations", 0)
set_input(modifier, group, "Fix Topology", 0)
set_input(modifier, group, "Smooth", 0.0)
evaluated = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())

4
tools/validate_moge_depth_plane.py
View File

@@ -8,6 +8,10 @@ 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,

82
tools/validate_moge_surface.py
View File

@@ -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)