diff options
Diffstat (limited to 'cnn_v3')
| -rw-r--r-- | cnn_v3/training/blender_export.py | 205 |
1 files changed, 113 insertions, 92 deletions
diff --git a/cnn_v3/training/blender_export.py b/cnn_v3/training/blender_export.py index 5d07c26..76180fd 100644 --- a/cnn_v3/training/blender_export.py +++ b/cnn_v3/training/blender_export.py @@ -25,8 +25,11 @@ G-buffer pass mapping: Alpha → transp. (0=opaque, 1=clear/transparent) """ +import os import sys import argparse +import shutil +import tempfile import bpy @@ -34,148 +37,166 @@ import bpy def parse_args(): # Blender passes its own argv; our args follow '--'. argv = sys.argv - if "--" in argv: - argv = argv[argv.index("--") + 1:] - else: - 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/frame_###", - help="Output path prefix (use ### for frame number padding). " - "Default: //renders/frame_###", - ) - 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)" + "--output", default="//renders/frame_###", + help="Output path prefix (use ### for frame number padding). Default: //renders/frame_###", ) + 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). " - "Use --view-layer without a value (or pass '?') to list available layers." + 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 - - # Render dimensions + 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 - - # Frame range (optional override) 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 - # Use Cycles for best multi-pass support - scene.render.engine = "CYCLES" - - # Enable required render passes on the selected view layer - available = list(scene.view_layers.keys()) - if getattr(args, "view_layer", None) in (None, "?"): - if args.view_layer == "?": - print("Available view layers:", available) - sys.exit(0) - vl = scene.view_layers[0] - else: - try: - vl = scene.view_layers[args.view_layer] - except KeyError: - print(f"ERROR: view layer '{args.view_layer}' not found.") - print(f"Available: {available}") - sys.exit(1) - 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; - # the compositor node below also taps it separately. + 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 - scene.use_nodes = True - tree = scene.node_tree - - # Clear all existing compositor nodes + tree = _get_compositor_tree(scene) tree.nodes.clear() - # Render Layers node (source of all passes) rl_node = tree.nodes.new("CompositorNodeRLayers") rl_node.location = (0, 0) - # File Output node — multi-layer EXR (all passes in one file) 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" - out_node.base_path = args.output + _set_output_path(out_node, args.output) - # Map each render pass socket to a named layer in the EXR. - # Slot order matters: the first slot is created by default; we rename it - # and add the rest. - pass_sockets = [ - ("Image", "Combined"), # beauty / target - ("Diffuse Color", "DiffCol"), # albedo - ("Normal", "Normal"), # world normals - ("Depth", "Z"), # view-space depth - ("Object Index", "IndexOB"), # object index - ("Shadow", "Shadow"), # shadow - ("Alpha", "Alpha"), # transparency / alpha - ] + # The node starts with one default slot; rename it first, then add the rest. + slots = getattr(out_node, "file_slots", None) or getattr(out_node, "file_output_items", 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) - # The node starts with one default slot; configure it first. - for i, (socket_name, layer_name) in enumerate(pass_sockets): - if i == 0: - # Rename the default slot - out_node.file_slots[0].path = layer_name - else: - out_node.file_slots.new(layer_name) + # 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 render layer socket to file output slot - src_socket = rl_node.outputs.get(socket_name) - dst_socket = out_node.inputs[i] - if src_socket: - tree.links.new(src_socket, dst_socket) + # 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: - print(f"[blender_export] WARNING: pass socket '{socket_name}' " - f"not found on Render Layers node. Skipping.") + # 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)) + print(" Layers: " + ", ".join(ln for _, ln, _ in PASS_SOCKETS)) def main(): args = parse_args() - configure_scene(args) + discard_dir = configure_scene(args) configure_compositor(args) - - # Trigger the render (only when running headless with -b) bpy.ops.render.render(animation=True) + shutil.rmtree(discard_dir, ignore_errors=True) print("[blender_export] Render complete.") |