feat: Add MoGe-2 support and enhance depth plane functionality
- Introduced new MoGe-2 depth plane group with geometry and material displacement options. - Updated existing depth plane group to improve background depth handling and subdivision settings. - Enhanced tests to validate MoGe-2 models and ensure proper integration with existing systems. - Modified installation environment to include new dependencies for MoGe-2 and utils3d. - Updated vendor lock file to include new artifacts for MoGe and utils3d. - Refactored download and model handling to accommodate new MoGe-2 models. - Improved validation scripts for depth plane assets to ensure compliance with new specifications.
This commit is contained in:
@@ -5,6 +5,7 @@ import bpy
|
||||
|
||||
GROUP_NAME = "O DA3 Surface"
|
||||
DEPTH_PLANE_GROUP_NAME = "O DA3 Depth Plane"
|
||||
MOGE_DEPTH_PLANE_GROUP_NAME = "O MoGe2 Depth Plane"
|
||||
MATERIAL_DISPLACEMENT_ATTRIBUTE = "o_material_displacement"
|
||||
BACKGROUND_DEPTH_ATTRIBUTE = "o_background_depth"
|
||||
DEPTH_SCALE_ATTRIBUTE = "o_depth_scale"
|
||||
@@ -689,7 +690,7 @@ def build_depth_plane_group():
|
||||
"INPUT",
|
||||
"NodeSocketInt",
|
||||
)
|
||||
subdivision.default_value = 5
|
||||
subdivision.default_value = 6
|
||||
subdivision.min_value = 0
|
||||
subdivision.max_value = 10
|
||||
displacement_strength = new_socket(
|
||||
@@ -708,7 +709,7 @@ def build_depth_plane_group():
|
||||
"INPUT",
|
||||
"NodeSocketFloat",
|
||||
)
|
||||
background_depth.default_value = -1.0
|
||||
background_depth.default_value = 1.0
|
||||
background_depth.min_value = -10000.0
|
||||
background_depth.max_value = 10000.0
|
||||
background_depth.subtype = "DISTANCE"
|
||||
@@ -786,6 +787,14 @@ def build_depth_plane_group():
|
||||
(350.0, 250.0),
|
||||
{"Background Depth"},
|
||||
)
|
||||
background_depth_value = nodes.new("ShaderNodeMath")
|
||||
background_depth_value.operation = "MULTIPLY"
|
||||
background_depth_value.inputs[1].default_value = -1.0
|
||||
background_depth_value.location = (550.0, 250.0)
|
||||
links.new(
|
||||
background_input.outputs["Background Depth"],
|
||||
background_depth_value.inputs[0],
|
||||
)
|
||||
displacement_input = new_group_input(
|
||||
nodes,
|
||||
(950.0, -150.0),
|
||||
@@ -906,7 +915,7 @@ def build_depth_plane_group():
|
||||
clamp_foreground.location = (600.0, 650.0)
|
||||
links.new(depth_texture.outputs["Color"], clamp_foreground.inputs[0])
|
||||
links.new(
|
||||
background_input.outputs["Background Depth"],
|
||||
background_depth_value.outputs[0],
|
||||
clamp_foreground.inputs[1],
|
||||
)
|
||||
foreground_mask = new_compare(
|
||||
@@ -925,7 +934,7 @@ def build_depth_plane_group():
|
||||
masked_depth.location = (800.0, 650.0)
|
||||
links.new(foreground_mask.outputs["Result"], masked_depth.inputs[0])
|
||||
links.new(
|
||||
background_input.outputs["Background Depth"],
|
||||
background_depth_value.outputs[0],
|
||||
masked_depth.inputs[2],
|
||||
)
|
||||
links.new(clamp_foreground.outputs[0], masked_depth.inputs[3])
|
||||
@@ -989,7 +998,7 @@ def build_depth_plane_group():
|
||||
(
|
||||
BACKGROUND_DEPTH_ATTRIBUTE,
|
||||
"FLOAT",
|
||||
attributes_input.outputs["Background Depth"],
|
||||
background_depth_value.outputs[0],
|
||||
),
|
||||
(
|
||||
DEPTH_SCALE_ATTRIBUTE,
|
||||
@@ -1020,11 +1029,225 @@ def build_depth_plane_group():
|
||||
return group
|
||||
|
||||
|
||||
def build_moge_depth_plane_group():
|
||||
group = bpy.data.node_groups.new(
|
||||
MOGE_DEPTH_PLANE_GROUP_NAME,
|
||||
"GeometryNodeTree",
|
||||
)
|
||||
group.is_modifier = True
|
||||
interface = group.interface
|
||||
new_socket(interface, "Geometry", "INPUT", "NodeSocketGeometry")
|
||||
displacement = new_socket(
|
||||
interface, "Displacement", "INPUT", "NodeSocketMenu"
|
||||
)
|
||||
subdivision = new_socket(
|
||||
interface, "Subdivision", "INPUT", "NodeSocketInt"
|
||||
)
|
||||
subdivision.default_value = 6
|
||||
subdivision.min_value = 0
|
||||
subdivision.max_value = 10
|
||||
depth_scale = new_socket(
|
||||
interface, "Depth Scale", "INPUT", "NodeSocketFloat"
|
||||
)
|
||||
depth_scale.default_value = 1.0
|
||||
depth_scale.structure_type = "SINGLE"
|
||||
background_depth = new_socket(
|
||||
interface, "Background Depth", "INPUT", "NodeSocketFloat"
|
||||
)
|
||||
background_depth.default_value = 1.0
|
||||
background_depth.min_value = -10000.0
|
||||
background_depth.max_value = 10000.0
|
||||
background_depth.subtype = "DISTANCE"
|
||||
background_depth.structure_type = "SINGLE"
|
||||
data_panel = interface.new_panel(name="Data", default_closed=True)
|
||||
depth_image = new_socket(
|
||||
interface, "Depth Image", "INPUT", "NodeSocketImage", data_panel
|
||||
)
|
||||
material = new_socket(
|
||||
interface, "Material", "INPUT", "NodeSocketMaterial", data_panel
|
||||
)
|
||||
for socket in (depth_image, material):
|
||||
socket.hide_in_modifier = True
|
||||
frame_sockets = []
|
||||
for name, socket_type, default in (
|
||||
("Frame", "NodeSocketInt", 1),
|
||||
("Frame Start", "NodeSocketInt", 1),
|
||||
("Frames", "NodeSocketInt", 1),
|
||||
("Cyclic", "NodeSocketBool", False),
|
||||
):
|
||||
socket = new_socket(interface, name, "INPUT", socket_type, data_panel)
|
||||
socket.default_value = default
|
||||
socket.hide_in_modifier = True
|
||||
frame_sockets.append(socket)
|
||||
new_socket(interface, "Geometry", "OUTPUT", "NodeSocketGeometry")
|
||||
|
||||
nodes = group.nodes
|
||||
links = group.links
|
||||
geometry_input = new_group_input(nodes, (-1600.0, 650.0), {"Geometry"})
|
||||
controls = new_group_input(
|
||||
nodes,
|
||||
(-1600.0, 150.0),
|
||||
{"Subdivision", "Displacement", "Depth Scale", "Background Depth"},
|
||||
)
|
||||
background_depth_value = nodes.new("ShaderNodeMath")
|
||||
background_depth_value.operation = "MULTIPLY"
|
||||
background_depth_value.inputs[1].default_value = -1.0
|
||||
background_depth_value.location = (-1350.0, -150.0)
|
||||
links.new(
|
||||
controls.outputs["Background Depth"],
|
||||
background_depth_value.inputs[0],
|
||||
)
|
||||
images = new_group_input(
|
||||
nodes, (-500.0, 1100.0), {"Depth Image"}
|
||||
)
|
||||
frames = new_group_input(
|
||||
nodes,
|
||||
(-500.0, 1500.0),
|
||||
{"Frame", "Frame Start", "Frames", "Cyclic"},
|
||||
)
|
||||
material_input = new_group_input(nodes, (900.0, 200.0), {"Material"})
|
||||
output = nodes.new("NodeGroupOutput")
|
||||
output.location = (1700.0, 650.0)
|
||||
|
||||
bounds = nodes.new("GeometryNodeBoundBox")
|
||||
bounds.location = (-1400.0, 650.0)
|
||||
links.new(geometry_input.outputs["Geometry"], bounds.inputs["Geometry"])
|
||||
dimensions = nodes.new("ShaderNodeVectorMath")
|
||||
dimensions.operation = "SUBTRACT"
|
||||
dimensions.location = (-1200.0, 650.0)
|
||||
links.new(bounds.outputs["Max"], dimensions.inputs[0])
|
||||
links.new(bounds.outputs["Min"], dimensions.inputs[1])
|
||||
separate = nodes.new("ShaderNodeSeparateXYZ")
|
||||
separate.location = (-1000.0, 650.0)
|
||||
links.new(dimensions.outputs["Vector"], separate.inputs[0])
|
||||
segments = nodes.new("ShaderNodeMath")
|
||||
segments.operation = "POWER"
|
||||
segments.inputs[0].default_value = 2.0
|
||||
segments.location = (-1200.0, 100.0)
|
||||
links.new(controls.outputs["Subdivision"], segments.inputs[1])
|
||||
vertices = nodes.new("ShaderNodeMath")
|
||||
vertices.operation = "ADD"
|
||||
vertices.inputs[1].default_value = 1.0
|
||||
vertices.location = (-1000.0, 100.0)
|
||||
links.new(segments.outputs[0], vertices.inputs[0])
|
||||
grid = nodes.new("GeometryNodeMeshGrid")
|
||||
grid.location = (-750.0, 650.0)
|
||||
links.new(separate.outputs["X"], grid.inputs["Size X"])
|
||||
links.new(separate.outputs["Y"], grid.inputs["Size Y"])
|
||||
links.new(vertices.outputs[0], grid.inputs["Vertices X"])
|
||||
links.new(vertices.outputs[0], grid.inputs["Vertices Y"])
|
||||
store_uv = nodes.new("GeometryNodeStoreNamedAttribute")
|
||||
store_uv.data_type = "FLOAT2"
|
||||
store_uv.domain = "CORNER"
|
||||
store_uv.inputs["Name"].default_value = "UVMap"
|
||||
store_uv.location = (-500.0, 650.0)
|
||||
links.new(grid.outputs["Mesh"], store_uv.inputs["Geometry"])
|
||||
links.new(grid.outputs["UV Map"], store_uv.inputs["Value"])
|
||||
|
||||
center = nodes.new("ShaderNodeVectorMath")
|
||||
center.operation = "ADD"
|
||||
center.location = (-1000.0, 900.0)
|
||||
links.new(bounds.outputs["Min"], center.inputs[0])
|
||||
links.new(bounds.outputs["Max"], center.inputs[1])
|
||||
half = nodes.new("ShaderNodeVectorMath")
|
||||
half.operation = "SCALE"
|
||||
half.inputs[3].default_value = 0.5
|
||||
half.location = (-750.0, 900.0)
|
||||
links.new(center.outputs["Vector"], half.inputs[0])
|
||||
transform = nodes.new("GeometryNodeTransform")
|
||||
transform.location = (-250.0, 650.0)
|
||||
links.new(store_uv.outputs["Geometry"], transform.inputs["Geometry"])
|
||||
links.new(half.outputs["Vector"], transform.inputs["Translation"])
|
||||
|
||||
depth_texture = nodes.new("GeometryNodeImageTexture")
|
||||
depth_texture.interpolation = "Linear"
|
||||
depth_texture.extension = "EXTEND"
|
||||
depth_texture.location = (-250.0, 1100.0)
|
||||
links.new(images.outputs["Depth Image"], depth_texture.inputs["Image"])
|
||||
links.new(grid.outputs["UV Map"], depth_texture.inputs["Vector"])
|
||||
frame = build_sequence_frame(
|
||||
nodes, links, frames.outputs["Frame"], frames, 1800.0
|
||||
)
|
||||
links.new(frame, depth_texture.inputs["Frame"])
|
||||
|
||||
clamp_depth = nodes.new("ShaderNodeMath")
|
||||
clamp_depth.operation = "MAXIMUM"
|
||||
clamp_depth.location = (0.0, 750.0)
|
||||
links.new(depth_texture.outputs["Color"], clamp_depth.inputs[0])
|
||||
links.new(background_depth_value.outputs[0], clamp_depth.inputs[1])
|
||||
displacement_menu = nodes.new("GeometryNodeMenuSwitch")
|
||||
displacement_menu.data_type = "FLOAT"
|
||||
displacement_menu.enum_items.clear()
|
||||
displacement_menu.enum_items.new("Geometry")
|
||||
displacement_menu.enum_items.new("Material")
|
||||
displacement_menu.inputs["Geometry"].default_value = 1.0
|
||||
displacement_menu.inputs["Material"].default_value = 0.0
|
||||
displacement_menu.location = (0.0, 300.0)
|
||||
links.new(controls.outputs["Displacement"], displacement_menu.inputs["Menu"])
|
||||
displacement.default_value = "Geometry"
|
||||
material_mode = nodes.new("ShaderNodeMath")
|
||||
material_mode.operation = "SUBTRACT"
|
||||
material_mode.inputs[0].default_value = 1.0
|
||||
material_mode.location = (250.0, 250.0)
|
||||
links.new(displacement_menu.outputs[0], material_mode.inputs[1])
|
||||
geometry_strength = nodes.new("ShaderNodeMath")
|
||||
geometry_strength.operation = "MULTIPLY"
|
||||
geometry_strength.location = (250.0, 500.0)
|
||||
links.new(displacement_menu.outputs[0], geometry_strength.inputs[0])
|
||||
links.new(controls.outputs["Depth Scale"], geometry_strength.inputs[1])
|
||||
scaled_depth = nodes.new("ShaderNodeMath")
|
||||
scaled_depth.operation = "MULTIPLY"
|
||||
scaled_depth.location = (250.0, 750.0)
|
||||
links.new(clamp_depth.outputs[0], scaled_depth.inputs[0])
|
||||
links.new(geometry_strength.outputs[0], scaled_depth.inputs[1])
|
||||
normal = nodes.new("GeometryNodeInputNormal")
|
||||
normal.location = (0.0, 500.0)
|
||||
offset = nodes.new("ShaderNodeVectorMath")
|
||||
offset.operation = "SCALE"
|
||||
offset.location = (250.0, 700.0)
|
||||
links.new(normal.outputs["Normal"], offset.inputs[0])
|
||||
links.new(scaled_depth.outputs[0], offset.inputs[3])
|
||||
set_position = nodes.new("GeometryNodeSetPosition")
|
||||
set_position.location = (500.0, 650.0)
|
||||
links.new(transform.outputs["Geometry"], set_position.inputs["Geometry"])
|
||||
links.new(offset.outputs["Vector"], set_position.inputs["Offset"])
|
||||
|
||||
set_material = nodes.new("GeometryNodeSetMaterial")
|
||||
set_material.location = (750.0, 650.0)
|
||||
links.new(set_position.outputs["Geometry"], set_material.inputs["Geometry"])
|
||||
links.new(material_input.outputs["Material"], set_material.inputs["Material"])
|
||||
smooth = nodes.new("GeometryNodeSetShadeSmooth")
|
||||
smooth.domain = "FACE"
|
||||
smooth.inputs["Shade Smooth"].default_value = True
|
||||
smooth.location = (1000.0, 650.0)
|
||||
links.new(set_material.outputs["Geometry"], smooth.inputs["Geometry"])
|
||||
geometry = smooth.outputs["Geometry"]
|
||||
for index, (name, data_type, value) in enumerate((
|
||||
(MATERIAL_DISPLACEMENT_ATTRIBUTE, "BOOLEAN", material_mode.outputs[0]),
|
||||
(BACKGROUND_DEPTH_ATTRIBUTE, "FLOAT", background_depth_value.outputs[0]),
|
||||
(DEPTH_SCALE_ATTRIBUTE, "FLOAT", controls.outputs["Depth Scale"]),
|
||||
)):
|
||||
store = nodes.new("GeometryNodeStoreNamedAttribute")
|
||||
store.data_type = data_type
|
||||
store.domain = "POINT"
|
||||
store.inputs["Name"].default_value = name
|
||||
store.location = (1200.0 + index * 200.0, 650.0)
|
||||
links.new(geometry, store.inputs["Geometry"])
|
||||
links.new(value, store.inputs["Value"])
|
||||
geometry = store.outputs["Geometry"]
|
||||
links.new(geometry, output.inputs["Geometry"])
|
||||
group.color_tag = "GEOMETRY"
|
||||
group.description = "Build a MoGe-2 depth plane with geometry or material displacement."
|
||||
group.use_fake_user = True
|
||||
return group
|
||||
|
||||
|
||||
def main():
|
||||
bpy.ops.wm.read_factory_settings(use_empty=True)
|
||||
bpy.context.preferences.filepaths.save_version = 0
|
||||
build_group()
|
||||
build_depth_plane_group()
|
||||
build_moge_depth_plane_group()
|
||||
OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
bpy.ops.wm.save_as_mainfile(filepath=str(OUTPUT_PATH), check_existing=False)
|
||||
|
||||
|
||||
@@ -29,7 +29,15 @@ def main():
|
||||
)
|
||||
assert background_depth.subtype == "DISTANCE"
|
||||
assert background_depth.structure_type == "SINGLE"
|
||||
assert background_depth.default_value == -1.0
|
||||
assert background_depth.default_value == 1.0
|
||||
subdivision = next(
|
||||
item
|
||||
for item in group.interface.items_tree
|
||||
if item.item_type == "SOCKET"
|
||||
and item.in_out == "INPUT"
|
||||
and item.name == "Subdivision"
|
||||
)
|
||||
assert subdivision.default_value == 6
|
||||
frame_input = next(
|
||||
item
|
||||
for item in group.interface.items_tree
|
||||
|
||||
79
tools/validate_moge_depth_plane.py
Normal file
79
tools/validate_moge_depth_plane.py
Normal file
@@ -0,0 +1,79 @@
|
||||
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_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)
|
||||
print("MoGe-2 depth plane validation passed")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user