Save each pass in a numpy array for a later usage in python

Contributing to Blender
1

Hello,

I am building an addon that can export passes to a numpy array. Unfortunately, it doesn’t seem like it is something one can do easily.

My current (and not working) solution is to get the pixel values from a viewer node bpy.data.images["Viewer Node"].pixels. While this work for one pass, when I try to change the input of the viewer node via a script, the pixel values does not update to the ones of the new pass (but the link is correctly created because when I switch back to the compositing tab, it is there).

My piece of code to change the viewer node’s input:

def pass_to_viewernode(pass_name: str):
    tree = bpy.context.scene.node_tree
    rl_node = tree.nodes.get("Render Layers")
    composite_node = tree.nodes.get("Viewer")
    render_pass = rl_node.outputs[pass_name]
    tree.links.new(render_pass, composite_node.inputs[0], True)

I found out that by manually moving the link back into the viewer node input will trigger an update of the pixel values, but it is not a behaviour I succeeded to reproduce via a script.

The other piece of code I am using to save the array. It will only save the same array n times, with n being the number of passes to save.

def get_pass_img(passname: str, n_channels: int, dtype: str):
    pass_to_viewernode(passname)

    bpy.data.images["Viewer Node"].reload()
    pixels = deepcopy(bpy.data.images["Viewer Node"].pixels)
    viewer = np.array(pixels)
    viewer = viewer.reshape((1024, 2048, -1))[:, :, :n_channels]

    viewer = viewer.astype(dtype)
    return viewer

bpy.ops.render.render(write_still=True)

rendered_passes = {}
for passname, pass_data in render_passes.items():
    pass_img = get_pass_img(passname, **pass_data)

    rendered_passes[passname] = pass_img

np.savez_compressed(savedir / str(image_id), **rendered_passes)

Do you have any workaround to trigger the update of the node via a script (I know that bpy.ops.render.render will do, but it will also rerender the entire frame which takes time considering the large number of assets) ?

Thank you for your time.

2

Well, there are some logican errors with your code, you can try below code and I hope it will resolve error.

import bpy
import numpy as np
from pathlib import Path

def render_passes_to_numpy(savedir, render_passes):
    bpy.context.scene.use_nodes = True
    tree = bpy.context.scene.node_tree

    for passname, pass_data in render_passes.items():
        pass_to_viewernode(passname)

        bpy.ops.render.render(write_still=True)

        pixels = bpy.data.images["Viewer Node"].pixels
        viewer = np.array(pixels)
        viewer = viewer.reshape((1024, 2048, -1))[:, :, :pass_data['n_channels']]

        viewer = viewer.astype(pass_data['dtype'])

        np.savez_compressed(savedir / f"{passname}.npz", data=viewer)

def pass_to_viewernode(pass_name):
    tree = bpy.context.scene.node_tree
    rl_node = tree.nodes.get("Render Layers")
    composite_node = tree.nodes.get("Viewer")

    for link in composite_node.inputs[0].links:
        tree.links.remove(link)

    render_pass = rl_node.outputs[pass_name]
    tree.links.new(render_pass, composite_node.inputs[0])

# Define the render passes and their properties
render_passes = {
    "Diffuse": {"n_channels": 3, "dtype": np.float32},
    "Normals": {"n_channels": 3, "dtype": np.float32},
    # Add more passes as needed
}

savedir = Path("/path/to/save/directory")
render_passes_to_numpy(savedir, render_passes)

Thanks

3

I made a few typo when translating the code for the forum. The code I had was functional (but wasn’t doing what I wanted).

function(**dict) is a way of using the pair key value of a dictionnary as argument of a function, it is not a logicial error ^^ (but I did make a typo to select only the number of channels I wanted)

The issue with this code is that it will render the entire image for each pass, which is not efficient at all (and takes a long time).

I was writing an answer about my alternative solution that I am using for now.

4

As it seems rather hard to do entirely through a script, I decided to use the OpenEXR format and to compile the library myself (rather cumbersome, that s why I wanted to avoid it).

This is the code I ended up with:

def exr_to_numpy(filepath, passdata: dict[str, dict]):
    exrfile = OpenEXR.InputFile(str(filepath))

    displaywindow = exrfile.header()['displayWindow']
    height = displaywindow.max.y + 1 - displaywindow.min.y
    width = displaywindow.max.x + 1 - displaywindow.min.x

    pass_arrays = {}
    for passname, pass_data in passdata.items():
        full_pass = get_exr_pass(exrfile, passname, **pass_data)

        full_pass = np.reshape(full_pass, (height, width, -1))

        pass_arrays[passname] = full_pass

    return pass_arrays


def get_exr_pass(exr_file, passname: str, channels: str, dtype: str = "float16") -> np.ndarray:
    """
    extract value of a pass from an exr file.

    Parameters
    ----------
    exr_file :
        The exr file object
    passname : str
        name of the pass as written in the exr file.
    channels : str
        name of the channels for the pass. It consists of a string of
        the following letters: V, X, Y, Z, R, G, B, A.
    dtype : str, optional
        dtype conversion for the channel, as float 32 may be unecessary,
        by default "float16"

    Returns
    -------
    np.ndarray
        1D array representing the pass.
    """
    channel_names = [f"{passname}.{channel}" for channel in channels]

    channels_raw: list[bytes]
    channels_raw = exr_file.channels(channel_names, Imath.PixelType(Imath.PixelType.FLOAT))

    channels = []
    for channel in channels_raw:
        # the type conversion must happen after the frombuffer loading in float32
        channel_values = np.frombuffer(channel, dtype=np.float32).astype(dtype)

        channels.append(channel_values)

    return np.stack(channels, 1)

render_passes = {
    "Depth": {"channels": "V", "dtype": "float16"},
    "Normal": {"channels": "XYZ", "dtype": "float16"},
    "DiffCol": {"channels": "RGB", "dtype": "float16"}
}

bpy.ops.render.render(write_still=True)

pass_arrays = exr_to_numpy(root_dir / "tmp0001.exr", render_passes)

np.savez_compressed(savedir / str(image_id), **pass_arrays)

The numpy compressed array will be a bit heavier than a compressed EXR file if you stick to float32. Otherwise, with half precision you ll have ~6times smaller files.