250 lines
8.7 KiB
Python
250 lines
8.7 KiB
Python
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" / "anyimage" / "wheels").glob("blendjob-*.whl"))),
|
|
)
|
|
|
|
from anyimage.operators.image import create_color_image
|
|
from anyimage.operators.material import (
|
|
MOGE_DEPTH_PLANE_NODE_GROUP_NAME,
|
|
MOGE_MATERIAL_NODE_GROUP_NAME,
|
|
get_modifier_input,
|
|
modifier_input_identifier,
|
|
)
|
|
from anyimage.operators.depth_plane.build import build_moge2_depth_plane
|
|
|
|
|
|
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()
|