path: root/cnn_v3/training/blender_export.py

"""
Blender export script for CNN v3 G-buffer training data.
Configures render passes and a compositor File Output node,
then renders the current scene to a multi-layer EXR.

Usage (headless):
    blender -b scene.blend -P blender_export.py -- --output /tmp/renders/

    # List available view layers in the blend file:
    blender -b scene.blend -P blender_export.py -- --view-layer ?

    # Use a specific view layer:
    blender -b scene.blend -P blender_export.py -- --output /tmp/renders/ --view-layer "MyLayer"

--output is the base directory for the compositor File Output node. Blender appends
the frame number to produce ####.exr files (e.g. /tmp/renders/0001.exr).
Use // as a prefix to resolve relative to the .blend file directory.

G-buffer pass mapping:
    Combined  → training target RGBA (beauty)
    DiffCol   → albedo.rgb  (pre-lighting material color)
    Normal    → normal.xy   (world-space, oct-encode in pack_blender_sample.py)
    Z         → depth       (view-space distance, normalize in pack step)
    IndexOB   → mat_id      (object index, u8 / 255)
    Shadow    → shadow      (1 = fully lit)
    Alpha     → transp.     (0=opaque, 1=clear/transparent)

Blender version notes:
    ≤4.x  Works: scene.node_tree compositor routes all per-pass render data.
    5.0.x BROKEN: the new post-process compositor (compositing_node_group) only
          receives the Combined image; DiffCol/Normal/Z/IndexOB/Shadow are empty.
          Use Blender 4.5 LTS until this is fixed upstream.
    5.x+  The compositing_node_group path is kept ready; once Blender routes
          per-pass buffers to CompositorNodeRLayers in the new compositor,
          this script will work without changes.
"""

import os
import sys
import argparse
import shutil

import bpy


def parse_args():
    # Blender passes its own argv; our args follow '--'.
    argv = sys.argv
    argv = argv[argv.index("--") + 1:] if "--" in argv else []
    parser = argparse.ArgumentParser(
        description="Configure Blender render passes and export multi-layer EXR."
    )
    parser.add_argument(
        "--output", default="//renders/",
        help="Base output directory for compositor File Output node. "
             "Blender appends frame number (e.g. 0001.exr). "
             "Use // for blend file directory. Default: //renders/",
    )
    parser.add_argument("--width",       type=int, default=640,  help="Render width in pixels (default: 640)")
    parser.add_argument("--height",      type=int, default=360,  help="Render height in pixels (default: 360)")
    parser.add_argument("--start-frame", type=int, default=None, help="First frame to render (default: scene start frame)")
    parser.add_argument("--end-frame",   type=int, default=None, help="Last frame to render (default: scene end frame)")
    parser.add_argument(
        "--view-layer", default=None, metavar="NAME",
        help="View layer name to use (default: first available). Pass '?' to list available layers.",
    )
    return parser.parse_args(argv)


def _die(msg, candidates=None):
    print(f"ERROR: {msg}")
    if candidates is not None:
        print("Candidate attributes:", candidates)
    sys.exit(1)


def _blender_version():
    """Return (major, minor) tuple, e.g. (4, 5) for Blender 4.5."""
    return bpy.app.version[:2]


def _is_blender5():
    """True on Blender 5.0+."""
    return _blender_version()[0] >= 5


def _resolve_view_layer(scene, name):
    """Return the requested view layer, defaulting to index 0. '?' lists and exits."""
    available = list(scene.view_layers.keys())
    if name == "?":
        print("Available view layers:", available)
        sys.exit(0)
    if name is None:
        return scene.view_layers[0]
    try:
        return scene.view_layers[name]
    except KeyError:
        _die(f"view layer '{name}' not found. Available: {available}")


def _get_compositor_tree(scene):
    """Return the compositor node tree, creating it if necessary.

    Blender ≤4.x: scene.node_tree (activated via scene.use_nodes = True).
    Blender 5.0+: scene.compositing_node_group.  NOTE: as of 5.0.x the new
    post-process compositor does not route per-pass render data to
    CompositorNodeRLayers — only the Combined image flows through.  The code
    below is correct and will work automatically once Blender fixes this.
    """
    if _is_blender5():
        if scene.compositing_node_group is None:
            scene.compositing_node_group = bpy.data.node_groups.new(
                name="Compositor", type="CompositorNodeTree"
            )
        return scene.compositing_node_group
    # Blender ≤4: use legacy scene.node_tree.
    if hasattr(scene, "use_nodes"):
        scene.use_nodes = True
    if hasattr(scene, "node_tree"):
        return scene.node_tree
    candidates = [a for a in dir(scene) if "node" in a.lower() or "compos" in a.lower()]
    _die("cannot find compositor node tree on scene.", candidates)


def _set_output_path(node, path):
    """Set the output path on a CompositorNodeOutputFile (API differs by version)."""
    for attr in ("base_path", "directory"):
        if hasattr(node, attr):
            setattr(node, attr, path)
            return
    candidates = [a for a in dir(node) if "path" in a.lower() or "dir" in a.lower() or "file" in a.lower()]
    _die("cannot find path attribute on CompositorNodeOutputFile.", candidates)


def configure_scene(args):
    """Configure render settings and enable G-buffer passes."""
    scene = bpy.context.scene

    major, minor = _blender_version()
    print(f"[blender_export] Blender {major}.{minor}")
    if _is_blender5():
        print("  WARNING: Blender 5.x detected. Per-pass compositor routing is broken "
              "in 5.0.x — only Combined will appear in the output EXR.")
        print("  Use Blender 4.5 LTS for training data generation.")

    scene.render.engine = "CYCLES"
    scene.render.resolution_x = args.width
    scene.render.resolution_y = args.height
    scene.render.resolution_percentage = 100
    if args.start_frame is not None:
        scene.frame_start = args.start_frame
    if args.end_frame is not None:
        scene.frame_end = args.end_frame

    vl = _resolve_view_layer(scene, args.view_layer)
    vl.use_pass_combined      = True   # beauty target
    vl.use_pass_diffuse_color = True   # albedo
    vl.use_pass_normal        = True   # world normals
    vl.use_pass_z             = True   # depth (Z)
    vl.use_pass_object_index  = True   # mat_id
    vl.use_pass_shadow        = True   # shadow
    # Alpha is available via the combined pass alpha channel.

    # Ensure the compositor runs during render (needed for File Output node).
    if hasattr(scene.render, "use_compositing"):
        scene.render.use_compositing = True

    # Redirect the render engine's own output to a discard directory; the real
    # output comes from the compositor File Output node written to args.output.
    discard_dir = os.path.join(os.path.abspath(args.output), "tmp")
    os.makedirs(discard_dir, exist_ok=True)
    scene.render.filepath = discard_dir + os.sep

    print(f"[blender_export] Render passes configured on ViewLayer '{vl.name}'.")
    print(f"  Resolution: {args.width}x{args.height}")
    print(f"  Frames: {scene.frame_start} – {scene.frame_end}")
    return discard_dir


# CompositorNodeRLayers output socket names (queried from Blender 5.0.1 console):
#   ng = bpy.data.node_groups.new("T", "CompositorNodeTree")
#   rl = ng.nodes.new("CompositorNodeRLayers")
#   print([s.name for s in rl.outputs])
#   bpy.data.node_groups.remove(ng)
#
#   ['Image', 'Alpha', 'Depth', 'Normal', 'UV', 'Vector', 'Position',
#    'Deprecated', 'Deprecated',
#    'Shadow', 'Ambient Occlusion',
#    'Deprecated', 'Deprecated', 'Deprecated',
#    'Object Index', 'Material Index', 'Mist', 'Emission', 'Environment',
#    'Diffuse Direct', 'Diffuse Indirect', 'Diffuse Color',
#    'Glossy Direct', 'Glossy Indirect', 'Glossy Color',
#    'Transmission Direct', 'Transmission Indirect', 'Transmission Color',
#    'Subsurface Direct', 'Subsurface Indirect', 'Subsurface Color']
#
# Table: (RenderLayer socket name, EXR layer name, Blender 5+ slot type)
PASS_SOCKETS = [
    ("Image",        "Combined", "RGBA"),   # beauty / target
    ("Diffuse Color","DiffCol",  "RGBA"),   # albedo
    ("Normal",       "Normal",   "VECTOR"), # world normals
    ("Depth",        "Z",        "FLOAT"),  # view-space depth
    ("Object Index", "IndexOB",  "FLOAT"),  # object index
    ("Shadow",       "Shadow",   "FLOAT"),  # shadow
    ("Alpha",        "Alpha",    "FLOAT"),  # transparency
]

# Fallback socket name aliases for older Blender versions.
_SOCKET_ALIASES = {
    "Shadow":        ["Shadow Catcher", "ShadowCatcher"],
    "Diffuse Color": ["DiffCol", "Diffuse"],
    "Depth":         ["Z"],
    "Object Index":  ["IndexOB"],
    "Alpha":         ["Alpha"],
}


def configure_compositor(args):
    """Build the compositor node graph: RenderLayers → FileOutput (OPEN_EXR_MULTILAYER)."""
    scene = bpy.context.scene
    vl = _resolve_view_layer(scene, args.view_layer)
    tree = _get_compositor_tree(scene)
    tree.nodes.clear()

    rl_node = tree.nodes.new("CompositorNodeRLayers")
    rl_node.location = (0, 0)
    # Bind to the correct scene + view layer so per-pass outputs are populated.
    if hasattr(rl_node, "scene"):
        rl_node.scene = scene
    if hasattr(rl_node, "layer"):
        rl_node.layer = vl.name

    out_node = tree.nodes.new("CompositorNodeOutputFile")
    out_node.location = (600, 0)
    out_node.format.file_format = "OPEN_EXR_MULTILAYER"
    out_node.format.exr_codec = "ZIP"
    _set_output_path(out_node, args.output)
    if hasattr(out_node, "file_name"):
        out_node.file_name = ""  # Blender 5: clear prefix so output is just ####.exr

    # Populate output slots.
    # file_output_items (Blender 5+) starts empty; add all slots.
    # file_slots (Blender ≤4) has one default slot; rename first, then append.
    slots = getattr(out_node, "file_output_items", None) \
         or getattr(out_node, "file_slots", None)
    if slots is None:
        candidates = [a for a in dir(out_node) if not a.startswith("__")]
        _die("cannot find slots attribute on CompositorNodeOutputFile.", candidates)

    if hasattr(out_node, "file_output_items"):
        for _, layer_name, sock_type in PASS_SOCKETS:
            slots.new(sock_type, layer_name)
    else:
        _, first_layer, _ = PASS_SOCKETS[0]
        slots[0].path = first_layer
        for _, layer_name, _ in PASS_SOCKETS[1:]:
            slots.new(layer_name)

    # Link each RenderLayer output socket to the corresponding FileOutput slot.
    for socket_name, layer_name, _ in PASS_SOCKETS:
        candidates = [socket_name] + _SOCKET_ALIASES.get(socket_name, [])
        src = next((s for s in rl_node.outputs if s.name in candidates), None)
        dst = out_node.inputs.get(layer_name)
        if src and dst:
            tree.links.new(src, dst)
        else:
            print(f"[blender_export] WARNING: cannot link '{socket_name}' → '{layer_name}'. Skipping.")

    print(f"[blender_export] Compositor configured. Output → {args.output}")
    print("  Layers: " + ", ".join(ln for _, ln, _ in PASS_SOCKETS))


def main():
    args = parse_args()
    discard_dir = configure_scene(args)
    configure_compositor(args)
    bpy.ops.render.render(animation=True)
    shutil.rmtree(discard_dir, ignore_errors=True)
    print("[blender_export] Render complete.")

    # Print the next-step pack command.
    pack_script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "pack_blender_sample.py")
    out_abs = os.path.abspath(args.output)
    print("\n[blender_export] Next step — pack EXRs into sample directories:")
    print(f"  for exr in {out_abs}/*.exr; do")
    print(f'    name=$(basename "${{exr%.exr}}")')
    print(f'    python3 {pack_script} --exr "$exr" --output {out_abs}/$name/')
    print( "  done")


if __name__ == "__main__":
    main()