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 slot name and frame number to each file (e.g. /tmp/renders/Combined0001.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      (invert: shadow=1 means fully lit)
    Alpha     → transp.     (0=opaque, 1=clear/transparent)
"""

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 slot name + frame number (e.g. Combined0001.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 _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.
    Handles API differences between Blender <=4.x and 5.x+."""
    if hasattr(scene, "compositing_node_group"):          # Blender 5.0+
        if scene.compositing_node_group is None:
            scene.compositing_node_group = bpy.data.node_groups.new(
                name="Compositor", type="CompositorNodeTree"
            )
        # In Blender 5.0+, assigning compositing_node_group activates compositing.
        return scene.compositing_node_group
    if hasattr(scene, "node_tree"):                        # Blender <= 4.x
        scene.use_nodes = True
        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):
    scene = bpy.context.scene
    scene.render.engine = "CYCLES"
    # Suppress default PNG output — all passes are captured via compositor File Output node.
    # Redirect to a tmp/ subdir next to --output; cleaned up after render (see main()).
    _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
    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 catcher
    # Alpha is available via the combined pass alpha channel.

    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


# Full CompositorNodeRLayers output socket list (Blender 5.0.1, queried via 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']
#
# Mapping used here  (socket name → EXR layer name → Blender 5+ type):
# (render-layer socket name, EXR layer name, Blender 5+ socket 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
]


def configure_compositor(args):
    scene = bpy.context.scene
    tree = _get_compositor_tree(scene)
    tree.nodes.clear()

    rl_node = tree.nodes.new("CompositorNodeRLayers")
    rl_node.location = (0, 0)

    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

    # The node starts with one default slot; rename it first, then add the rest.
    slots = getattr(out_node, "file_output_items", None)
    if slots is None:
        slots = 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)

    # Populate slots. file_output_items (Blender 5+) starts empty;
    # file_slots (older) has one default slot that must be renamed before adding more.
    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)

    # Socket name aliases across Blender versions.
    SOCKET_ALIASES = {
        "Shadow":        ["Shadow Catcher", "ShadowCatcher"],
        "Diffuse Color": ["DiffCol", "Diffuse"],
        "Depth":         ["Z"],
        "Object Index":  ["IndexOB"],
        "Alpha":         ["Alpha"],
    }

    # Link each render-layer output to the file output slot by name.
    for socket_name, layer_name, _ in PASS_SOCKETS:
        candidates = [socket_name] + SOCKET_ALIASES.get(socket_name, [])
        src = next((rl_node.outputs.get(n) for n in candidates if rl_node.outputs.get(n)), 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 (HOW_TO_CNN.md §"Batch packing").
    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()