diff --git a/infinigen/core/nodes/node_info.py b/infinigen/core/nodes/node_info.py index 11755447..8144a0b4 100644 --- a/infinigen/core/nodes/node_info.py +++ b/infinigen/core/nodes/node_info.py @@ -177,6 +177,7 @@ class Nodes: WorldOutput = "ShaderNodeOutputWorld" Composite = "CompositorNodeComposite" Viewer = "CompositorNodeViewer" + CompositorMixRGB = "CompositorNodeMixRGB" # Point DistributePointsOnFaces = "GeometryNodeDistributePointsOnFaces" diff --git a/infinigen/core/rendering/post_render.py b/infinigen/core/rendering/post_render.py index d0b0b79a..9a1f2aee 100644 --- a/infinigen/core/rendering/post_render.py +++ b/infinigen/core/rendering/post_render.py @@ -8,6 +8,8 @@ import logging import os +import OpenEXR + # ruff: noqa: E402 os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1" # This must be done BEFORE import cv2. @@ -30,8 +32,21 @@ def load_exr(path): load_flow = load_exr -def load_depth(p): - return load_exr(p)[..., 0] +def load_single_channel(p): + file = OpenEXR.InputFile(str(p)) + channel, channel_type = next(iter(file.header()["channels"].items())) + match str(channel_type.type): + case "FLOAT": + np_type = np.float32 + case _: + np_type = np.uint8 + data = np.frombuffer(file.channel(channel, channel_type.type), np_type) + dw = file.header()["dataWindow"] + sz = (dw.max.x - dw.min.x + 1, dw.max.y - dw.min.y + 1) + return data.reshape(sz) + + +load_depth = load_single_channel def load_normals(p): @@ -39,7 +54,7 @@ def load_normals(p): def load_seg_mask(p): - return load_exr(p)[..., 2].astype(np.int64) + return load_single_channel(p).astype(np.int64) def load_uniq_inst(p): diff --git a/infinigen/core/rendering/render.py b/infinigen/core/rendering/render.py index 75bfa9da..43cc44d8 100644 --- a/infinigen/core/rendering/render.py +++ b/infinigen/core/rendering/render.py @@ -192,14 +192,23 @@ def configure_compositor_output( slot_input = file_output_node.file_slots.new(socket_name) render_socket = render_layers.outputs[socket_name] - if viewlayer_pass == "vector": - separate_color = nw.new_node(Nodes.CompSeparateColor, [render_socket]) - comnbine_color = nw.new_node( - Nodes.CompCombineColor, [0, (separate_color, 3), (separate_color, 2), 0] - ) - nw.links.new(comnbine_color.outputs[0], slot_input) - else: - nw.links.new(render_socket, slot_input) + match viewlayer_pass: + case "vector": + separate_color = nw.new_node(Nodes.CompSeparateColor, [render_socket]) + comnbine_color = nw.new_node( + Nodes.CompCombineColor, + [0, (separate_color, 3), (separate_color, 2), 0], + ) + nw.links.new(comnbine_color.outputs[0], slot_input) + case "normal": + color = nw.new_node( + Nodes.CompositorMixRGB, + [None, render_socket, (0, 0, 0, 0)], + attrs={"blend_type": "ADD"}, + ).outputs[0] + nw.links.new(color, slot_input) + case _: + nw.links.new(render_socket, slot_input) file_slot_list.append(file_output_node.file_slots[slot_input.name]) slot_input = default_file_output_node.file_slots["Image"] diff --git a/infinigen/core/util/blender.py b/infinigen/core/util/blender.py index c22f88b5..86a3a661 100644 --- a/infinigen/core/util/blender.py +++ b/infinigen/core/util/blender.py @@ -22,7 +22,6 @@ from infinigen.core.nodes.node_info import DATATYPE_DIMS, DATATYPE_FIELDS -from ..nodes.node_wrangler import ng_inputs from . import math as mutil from .logging import Suppress @@ -556,7 +555,11 @@ def get_camera_res(): def set_geomod_inputs(mod, inputs: dict): assert mod.type == "NODES" for k, v in inputs.items(): - inputs = ng_inputs(mod.node_group) + inputs = { + s.name: s + for s in mod.node_group.interface.items_tree + if s.in_out == "INPUT" + } if k not in inputs: raise KeyError(f"Couldnt find {k=} in {mod.node_group.inputs.keys()=}") soc = inputs[k] diff --git a/infinigen/terrain/surface_kernel/kernelizer.py b/infinigen/terrain/surface_kernel/kernelizer.py index 2db9dc7d..dad7a2ef 100644 --- a/infinigen/terrain/surface_kernel/kernelizer.py +++ b/infinigen/terrain/surface_kernel/kernelizer.py @@ -13,7 +13,6 @@ from infinigen.terrain.utils import ( NODE_ATTRS_AVAILABLE, NODE_FUNCTIONS, - SOCKETTYPE_KERNEL, KernelDataType, Nodes, SocketType, @@ -29,6 +28,7 @@ value_string, var_list, ) +from infinigen.terrain.utils.kernelizer_util import SOCKETTYPE_KERNEL, SOCKETTYPES functional_nodes = [ Nodes.SetPosition, @@ -50,9 +50,10 @@ class Kernelizer: def get_inputs(self, node_tree): inputs = OrderedDict() for node_input in ng_inputs(node_tree).values(): - if node_input.type != SocketType.Geometry: - assert node_input.type != SocketType.Image - inputs[node_input.identifier] = SOCKETTYPE_KERNEL[node_input.type] + socket_type = SOCKETTYPES[node_input.socket_type] + if socket_type != SocketType.Geometry: + assert socket_type != SocketType.Image + inputs[node_input.identifier] = SOCKETTYPE_KERNEL[socket_type] return inputs def get_output(self, node_tree): @@ -61,8 +62,9 @@ def get_output(self, node_tree): if node.bl_idname == Nodes.SetPosition: outputs[Vars.Offset] = KernelDataType.float3 for node_output in ng_outputs(node_tree).values(): - if node_output.type != SocketType.Geometry: - outputs[node_output.identifier] = SOCKETTYPE_KERNEL[node_output.type] + socket_type = SOCKETTYPES[node_output.socket_type] + if socket_type != SocketType.Geometry: + outputs[node_output.identifier] = SOCKETTYPE_KERNEL[socket_type] return outputs def regularize(self, node_tree): @@ -418,8 +420,7 @@ def __call__(self, modifier): node_tree, collective_style=True ) for nodeoutput in ng_outputs(node_tree).values(): - id = nodeoutput.identifier - if id != "Output_1": # not Geometry + if nodeoutput.socket_type != "NodeSocketGeometry": code = re.sub(rf"\b{id}\b", modifier[f"{id}_attribute_name"], code) outputs[modifier[f"{id}_attribute_name"]] = outputs.pop(id) return code, imp_inputs, outputs diff --git a/infinigen/terrain/utils/kernelizer_util.py b/infinigen/terrain/utils/kernelizer_util.py index 72545c1a..f29695a6 100644 --- a/infinigen/terrain/utils/kernelizer_util.py +++ b/infinigen/terrain/utils/kernelizer_util.py @@ -152,6 +152,14 @@ class KernelDataType: # BOOL todo when necessary } +SOCKETTYPES = { + "NodeSocketFloat": SocketType.Value, + "NodeSocketVector": SocketType.Vector, + "NodeSocketInt": SocketType.Int, + "NodeSocketColor": SocketType.RGBA, + "NodeSocketImage": SocketType.Image, + "NodeSocketGeometry": SocketType.Geometry, +} NODE_FUNCTIONS = { Nodes.WaveTexture: "node_shader_tex_wave",