summaryrefslogtreecommitdiff
path: root/cnn_v3/training/blender_export.py
blob: 345f395778e89e7c503cdeaf54b9537f9882ecbb (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
"""
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 tempfile

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"
            )
        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 temp dir; files are cleaned up after render (see main()).
    _discard_dir = tempfile.mkdtemp(prefix="blender_export_discard_")
    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


# (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)

    # Link each render-layer output to the file output slot by name.
    for socket_name, layer_name, _ in PASS_SOCKETS:
        src = rl_node.outputs.get(socket_name)
        dst = out_node.inputs.get(layer_name)
        if src and dst:
            tree.links.new(src, dst)
        else:
            # TODO: socket names vary by Blender version / blend file (e.g. "Shadow" may be
            # "Shadow Catcher" or absent). Add a fallback name-matching step if needed.
            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.")


if __name__ == "__main__":
    main()