import sys import tempfile from pathlib import Path from types import SimpleNamespace from unittest.mock import patch 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"))), ) import snapform.blender.depth_plane as depth_plane from snapform.blender.depth_plane import ( MOGE_DEPTH_PLANE_NODE_GROUP_NAME, MOGE_MATERIAL_NODE_GROUP_NAME, build_moge2_depth_plane, create_color_image, get_modifier_input, modifier_input_identifier, ) def validate_single_frame_source_types(): with tempfile.TemporaryDirectory() as temporary: directory = Path(temporary) color_path = directory / "color.0001.png" seed = bpy.data.images.new("Color Seed", width=2, height=2, alpha=True) seed.pixels[:] = [0.2, 0.4, 0.8, 1.0] * 4 seed.filepath_raw = str(color_path) seed.file_format = "PNG" seed.save() bpy.data.images.remove(seed, do_unlink=True) still_source_image = bpy.data.images.load( str(color_path), check_existing=False ) still_source = bpy.data.objects.new("Still Source", None) still_source.empty_display_type = "IMAGE" still_source.data = still_source_image bpy.context.collection.objects.link(still_source) still_color = create_color_image( still_source, directory, assign_to_source=False ) assert still_color.source == "FILE" animated_source = SimpleNamespace( data=SimpleNamespace( source="SEQUENCE", name="Sequence Source.png", use_fake_user=False, ), image_user=None, ) sequence_color = create_color_image( animated_source, directory, assign_to_source=False ) assert sequence_color.source == "SEQUENCE" depth = np.ones((1, 2, 2), dtype=np.float32) mask = np.ones((1, 2, 2), dtype=bool) intrinsics = np.array( [[[1.0, 0.0, 0.5], [0.0, 1.0, 0.5], [0.0, 0.0, 1.0]]], dtype=np.float32, ) path = directory / "single-frame-results.npz" np.savez(path, depth=depth, mask=mask, intrinsics=intrinsics) with patch.object(depth_plane, "is_animated_image", return_value=True): plane = build_moge2_depth_plane( bpy.context, still_source, path, delete_source=False, color_image=sequence_color, data_frames_directory=directory, ) textures = [ node.image for node in plane.data.materials[0].node_tree.nodes if node.bl_idname == "ShaderNodeTexImage" ] assert sequence_color in textures depth_images = [image for image in textures if image != sequence_color] assert len(depth_images) == 1 assert depth_images[0].source == "SEQUENCE" 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, ) material_image = image.copy() material_image.name = "MoGe Test.color" plane = build_moge2_depth_plane( bpy.context, source, path, delete_source=True, color_image=material_image, ) 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 assert sockets["Color Image"].hide_in_modifier color_identifier = modifier_input_identifier(geometry_group, "Color Image") assert get_modifier_input(plane.modifiers[0], color_identifier) == material_image 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) static_depth = next( node.image for node in image_textures if node.image.name == "MoGe Test.depth" ) assert static_depth.source != "SEQUENCE" geometry_textures = [ node for node in geometry_group.nodes if node.bl_idname == "GeometryNodeImageTexture" ] assert len(geometry_textures) == 2 assert all(node.extension == "EXTEND" for node in geometry_textures) color_texture = next( node for node in geometry_textures if any( link.to_node == node and link.to_socket == node.inputs["Image"] and link.from_socket.name == "Color Image" for link in geometry_group.links ) ) depth_texture = next(node for node in geometry_textures if node != color_texture) assert color_texture.interpolation == "Closest" assert depth_texture.interpolation == "Linear" color_frame_link = next( link for link in geometry_group.links if link.to_node == color_texture and link.to_socket == color_texture.inputs["Frame"] ) depth_frame_link = next( link for link in geometry_group.links if link.to_node == depth_texture and link.to_socket == depth_texture.inputs["Frame"] ) assert color_frame_link.from_socket == depth_frame_link.from_socket alpha_compare = next( node for node in geometry_group.nodes if node.bl_idname == "FunctionNodeCompare" and node.operation == "GREATER_THAN" and any( link.from_node == color_texture and link.from_socket == color_texture.outputs["Alpha"] and link.to_node == node for link in geometry_group.links ) ) masked_depth = next( node for node in geometry_group.nodes if node.bl_idname == "ShaderNodeMix" and any( link.from_node == alpha_compare and link.to_node == node for link in geometry_group.links ) ) assert masked_depth.data_type == "FLOAT" masked_sources = { link.from_node for link in geometry_group.links if link.to_node == masked_depth } assert alpha_compare in masked_sources assert any( node.bl_idname == "ShaderNodeMath" and node.operation == "MAXIMUM" for node in masked_sources ) assert any( node.bl_idname == "ShaderNodeMath" and node.operation == "MULTIPLY" and node.inputs[1].default_value == -1.0 for node in masked_sources ) assert any( link.from_node == masked_depth and link.to_node.bl_idname == "ShaderNodeMath" and link.to_node.operation == "MULTIPLY" for link in geometry_group.links ) validate_single_frame_source_types() print("MoGe-2 depth plane validation passed") if __name__ == "__main__": main()