diff --git a/README.md b/README.md index e99d54c..d6d040c 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,7 @@ To use a procedural pipeline inside Blender GUI, follow these steps: #### Generate Assets using Standalone Scripts -Procedural assets can also be generated using standalone scripts that require close integration with Blender's Python API `bpy`. You can either write your own script that performs the steps outlined above or use the [procgen_assets.py](https://github.com/AndrejOrsula/space_robotics_bench/blob/main/scripts/blender/procgen_assets.py) as a starting point. +Procedural assets can also be generated using standalone scripts that require close integration with Blender's Python API `bpy`. You can either write your own script that performs the steps outlined above or use the [`procgen_assets.py`](scripts/blender/procgen_assets.py) as a starting point. This script provides an automated pipeline for generating procedural 3D models with baked textures into one of many supported formats (ABC, FBX, GLTF, OBJ, PLY, SDF, STL, USD). ## Attributions diff --git a/scripts/blender/__init__.py b/scripts/blender/__init__.py new file mode 100755 index 0000000..636498f --- /dev/null +++ b/scripts/blender/__init__.py @@ -0,0 +1,23 @@ +#!/usr/bin/env -S blender --python-exit-code 1 --python +""" +Script for setting up Blender preferences. + +- Use Cycles with CUDA GPU +- Enable Node Wrangler addon +""" + +import bpy + +# Use Cycles with GPU +bpy.context.preferences.addons["cycles"].preferences.compute_device_type = "CUDA" +bpy.context.preferences.addons["cycles"].preferences.get_devices() +for device in bpy.context.preferences.addons["cycles"].preferences.devices: + device.use = device.type == "CUDA" +bpy.context.scene.render.engine = "CYCLES" +bpy.context.scene.cycles.device = "GPU" + +# Enable Node Wrangler +bpy.ops.preferences.addon_enable(module="node_wrangler") + +# Save preferences +bpy.ops.wm.save_userpref() diff --git a/scripts/blender/procgen_assets.py b/scripts/blender/procgen_assets.py new file mode 100755 index 0000000..7b3f785 --- /dev/null +++ b/scripts/blender/procgen_assets.py @@ -0,0 +1,1293 @@ +#!/usr/bin/env -S blender --factory-startup --background --offline-mode --enable-autoexec --python-exit-code 1 --python +""" +Script for automated procedural asset generation using Blender that revolves around its rich +node-based system for geometry (Geometry Nodes) and materials (Shader Nodes). + +Overview: + The requested node trees are constructed via scripts defined by `--autorun_scripts`. A sequence of + Geometry Nodes moodifiers is then applied to a prototype object that is duplicated for each generated + variant. Once the geometry is finalized, a procedural material is applied and baked into PBR textures + before exporting the final model. + +Example (manual invocation is not recommended): + blender --python procgen_assets.py -- \ + --autorun_scripts path/to/nodes_0.py path/to/nodes_1.py ... \ + --geometry_nodes '{"NodeModifierName": {"input_name": input_value, ...}, ...}' \ + --material MaterialName \ + --texture_resolution 4096 \ + --num_assets 10 \ + --ext usdz \ + --outdir path/to/output/directory +""" + +from __future__ import annotations + +import argparse +import contextlib +import enum +import io +import json +import math +import os +import re +import sys +import time +from copy import deepcopy +from os import path +from pathlib import Path +from typing import Any, Dict, List, Literal, Optional, TextIO, Tuple, Union +from xml.dom import minidom as xml_dom +from xml.etree import ElementTree as xml_et + +import bpy +import mathutils + + +def main(**kwargs): + print_bpy(f"[INFO]: Starting procedural asset generation with kwargs: {kwargs}") + verify_requirements() + ProceduralGenerator.generate(**kwargs) + + +class ProceduralGenerator: + """ + Generator of procedural models using Blender. + """ + + @classmethod + def generate( + cls, + *, + ## Input + autorun_scripts: List[str], + ## Output + outdir: str, + name: str, + ext: str, + overwrite_min_age: int, + ## Generator + seed: int, + num_assets: int, + ## Export + export_kwargs: Dict[str, Any], + render_thumbnail: bool, + thumbnail_resolution: int, + ## Geometry + geometry_nodes: Dict[str, Dict[str, Any]], + decimate_angle_limit: Optional[float], + decimate_face_count: Optional[int], + ## Material + material: Optional[str], + texture_resolution: int, + render_samples: int, + ): + """ + Entrypoint method for generating a set of procedural models. + """ + + # Preprocess the input arguments + outdir = Path(outdir).absolute().as_posix() + if not ext.startswith("."): + ext = f".{ext}" + + # Create the output directory + os.makedirs(name=outdir, exist_ok=True) + + # Reset to factory settings with an empty scene + bpy.ops.wm.read_factory_settings(use_empty=True) + + # Autorun all input scripts + for script in autorun_scripts: + print_bpy(f"[INFO]: Running script: {script}") + bpy.ops.script.python_file_run(filepath=script) + + # Extract a map of the input socket mapping for all loaded node groups + node_group_input_socket_maps = cls.extract_aliased_input_socket_id_mappings() + + # Create an empty mesh object and treat it as a prototype that will be duplicated for each processed seed + bpy.ops.object.add(type="MESH") + proto_obj: bpy.types.Object = bpy.context.active_object + proto_obj.name = name + proto_obj.data.name = name + proto_obj.hide_render = True + + # Apply the requested list of geometry nodes to the prototype object + cls.apply_geometry_nodes_modifiers( + obj=proto_obj, + geometry_nodes=geometry_nodes, + node_group_input_socket_maps=node_group_input_socket_maps, + ) + + # Preheat the oven for baking the material into PBR textures + if material: + ProceduralGenerator.Baker.preheat_oven(render_samples=render_samples) + + # Prepare the scene for rendering the thumbnail + if render_thumbnail: + ProceduralGenerator.Exporter.prepare_renderer() + + # Generate models over the specified range + for current_seed in range(seed, seed + num_assets): + # Form the output filepath + filepath = path.join(outdir, f"{name}{current_seed}") + + # Skip generation if the file already exists and is too recent + if path.exists(filepath) and ( + overwrite_min_age < 0 + or (overwrite_min_age > time.time() - path.getmtime(filepath)) + ): + print_bpy( + f"[INFO]: Skipping generation of '{filepath}' because it was generated in the last {overwrite_min_age} seconds" + ) + continue + + # Duplicate the prototype object, rename it, select it and mark it as the active object + obj = cls.duplicate_object(obj=proto_obj, seed=current_seed) + + # Update the random seed of all geometry nodes modifiers that have a seed input socket + geometry_nodes_modifiers = [ + modifier + for modifier in obj.modifiers.values() + if modifier.type == "NODES" + ] + for modifier in geometry_nodes_modifiers: + if seed_id := node_group_input_socket_maps[ + modifier.node_group.name + ].get("seed"): + modifier[seed_id] = current_seed + + # Apply changes via mesh update + obj.data.update() + + # Apply all modifiers + for modifier in obj.modifiers: + bpy.ops.object.modifier_apply(modifier=modifier.name) + + # If specified, bake the material into PBR textures + if material: + obj.hide_render = False + obj.data.materials.clear() + obj.data.materials.append(bpy.data.materials.get(material)) + ProceduralGenerator.Baker.bake_into_pbr_material( + obj=obj, texture_resolution=texture_resolution + ) + + # Decimate the mesh if necessary + if decimate_angle_limit: + bpy.ops.object.modifier_add(type="DECIMATE") + obj.modifiers["Decimate"].decimate_type = "DISSOLVE" + obj.modifiers["Decimate"].angle_limit = decimate_angle_limit + + if decimate_face_count: + # Decimate the mesh + bpy.ops.object.modifier_add(type="DECIMATE") + obj.modifiers["Decimate"].ratio = decimate_face_count / len( + obj.data.polygons + ) + bpy.ops.object.modifier_apply(modifier="Decimate") + + # Export the model + cls.Exporter.export( + filepath=filepath, + ext=ext, + render_thumbnail=render_thumbnail, + thumbnail_resolution=thumbnail_resolution, + **export_kwargs, + ) + print_bpy( + f"[LOG]: Generated asset #{current_seed} ({current_seed-seed+1}/{num_assets}): {filepath}" + ) + + # Update the viewport to keep track of progress + if not bpy.app.background: + bpy.ops.wm.redraw_timer(type="DRAW_WIN_SWAP", iterations=1) + + # Remove the generated object + bpy.data.objects.remove(obj) + + print_bpy("[INFO]: Generation completed") + + @staticmethod + def extract_aliased_input_socket_id_mappings() -> Dict[str, Dict[str, str]]: + def _extract_input_socket_id_mapping( + node_group: bpy.types.NodeTree, + ) -> Dict[str, str]: + return { + canonicalize_str(item.name): item.identifier + for item in node_group.interface.items_tree.values() + if item.item_type == "SOCKET" and item.in_out == "INPUT" + } + + # Extract a map of the input socket mapping for all node groups + node_group_input_socket_maps = { + node_group_name: _extract_input_socket_id_mapping(node_group) + for node_group_name, node_group in bpy.data.node_groups.items() + } + + # Rename common aliases for convenience + COMMON_ALIASES: Dict[str, List[str]] = { + "seed": [ + "pseodorandomseed", + "randomseed", + "rng", + ], + "detail": [ + "detaillevel", + "detailobject", + "levelofdetail", + "subdivisionlevel", + "subdivisions", + "subdivlevel", + ], + } + for node_group_input_socket_map in node_group_input_socket_maps.values(): + for target, possible_alias in COMMON_ALIASES.items(): + original_alias: Optional[str] = ( + target if target in node_group_input_socket_map.keys() else None + ) + for key in node_group_input_socket_map.keys(): + if key in possible_alias: + if original_alias is not None: + raise ValueError( + "Ambiguous name of the input socket '{target}' (canonicalized): '{original_alias}', '{key}'" + ) + original_alias = key + if original_alias is not None and original_alias != target: + node_group_input_socket_map[target] = node_group_input_socket_map[ + original_alias + ] + + return node_group_input_socket_maps + + @staticmethod + def apply_geometry_nodes_modifiers( + obj: bpy.types.Object, + *, + geometry_nodes: Dict[str, Dict[str, Any]], + node_group_input_socket_maps: Dict[str, Dict[str, str]], + ): + for i, (node_group_name, node_group_inputs) in enumerate( + geometry_nodes.items() + ): + # Create a new nodes modifier + modifier: bpy.types.NodesModifier = obj.modifiers.new( + name=f"node{i}", type="NODES" + ) + + # Assign the requested node group + if node_group := bpy.data.node_groups.get(node_group_name): + modifier.node_group = node_group + else: + raise ValueError( + f"Node group '{node_group_name}' not found in the list of available groups: {bpy.data.node_groups.keys()}" + ) + + # Set inputs accordingly + for key, value in node_group_inputs.items(): + socket_id = node_group_input_socket_maps[node_group_name][ + canonicalize_str(key) + ] + modifier[socket_id] = value + + # Apply changes via mesh update + obj.data.update() + + @staticmethod + def duplicate_object(obj: bpy.types.Object, seed: int) -> bpy.types.Object: + # Select the object to duplicate + bpy.ops.object.select_all(action="DESELECT") + obj.select_set(True) + bpy.context.view_layer.objects.active = obj + + # Duplicate the object + bpy.ops.object.duplicate() + + # Get the new duplicated object (it will be the active object) + new_obj = bpy.context.active_object + + # Rename the new object based on the seed + new_obj.name = f"{obj.name}_{seed}" + new_obj.data.name = f"{obj.data.name}_{seed}" + + # Select the new object and make it the active object + bpy.ops.object.select_all(action="DESELECT") + new_obj.select_set(True) + bpy.context.view_layer.objects.active = new_obj + + return new_obj + + class Recipe(enum.Enum): + ALBEDO = enum.auto() + METALLIC = enum.auto() + SPECULAR = enum.auto() + ROUGHNESS = enum.auto() + NORMAL = enum.auto() + + @staticmethod + def enabled_recipes(): + return ( + ProceduralGenerator.Recipe.ALBEDO, + ProceduralGenerator.Recipe.METALLIC, + # ProceduralGenerator.Recipe.SPECULAR, + ProceduralGenerator.Recipe.ROUGHNESS, + ProceduralGenerator.Recipe.NORMAL, + ) + + @property + def bake_type(self): + match self: + case self.ALBEDO: + return "DIFFUSE" + case self.METALLIC: + return "EMIT" + case self.SPECULAR: + return "GLOSSY" + case _: + return self.name + + @property + def color_space(self): + return "sRGB" if self.ALBEDO == self else "Non-Color" + + @property + def shader_socket_name(self): + match self: + case self.ALBEDO: + return "Base Color" + case self.METALLIC: + return "Metallic" + case self.SPECULAR: + return "Specular Tint" + case self.ROUGHNESS: + return "Roughness" + case self.NORMAL: + return "Normal" + + def prep(self, material: bpy.types.Material) -> Dict[str, Any]: + match self: + # If the shader node has a "Metallic" input, it needs to be disabled for baking + case self.ALBEDO: + # Get the output material node + output_material_node = [ + node + for node in material.node_tree.nodes + if node.type == "OUTPUT_MATERIAL" + ][0] + + # Get surface shader socket + shader_node_socket = output_material_node.inputs["Surface"] + if not shader_node_socket.is_linked: + return {} + + # Get the shader node connected to the socket + shader_node = shader_node_socket.links[0].from_node + + # No need to do anything if the shader node does not have a "Metallic" input + if "Metallic" not in shader_node.inputs: + return {} + + metallic_socket = shader_node.inputs["Metallic"] + ingredients = {} + + # If the metallic input is linked, store the original link source and temporarily disconnect it + if metallic_socket.is_linked: + metallic_socket_links = metallic_socket.links[0] + ingredients["orig_metallic_from_socket"] = ( + metallic_socket_links.from_socket + ) + material.node_tree.links.remove(metallic_socket_links) + + # Always store the original default value and set it to 0.0 + ingredients["orig_metallic_default_value"] = ( + metallic_socket.default_value + ) + metallic_socket.default_value = 0.0 + + return ingredients + + # Render the metallic input as emission originating from the surface shader's metallic input + case self.METALLIC: + # Get the output material node + output_material_node = [ + node + for node in material.node_tree.nodes + if node.type == "OUTPUT_MATERIAL" + ][0] + + # Get surface shader socket + shader_node_socket = output_material_node.inputs["Surface"] + if not shader_node_socket.is_linked: + return {} + + # Get the shader link + shader_link = shader_node_socket.links[0] + + # Get the shader node connected to the socket + shader_node = shader_link.from_node + + # Store the original link source and temporarily disconnect it + ingredients = { + "orig_surface_shader_source": shader_link.from_socket + } + material.node_tree.links.remove(shader_link) + + # No need to do anything if the shader node does not have a "Metallic" input + orig_metallic_default_value = 0.0 + with_emissive_rgb_node = False + if "Metallic" in shader_node.inputs: + metallic_socket = shader_node.inputs["Metallic"] + + # If the metallic input is linked, store the original link source and temporarily disconnect it + if metallic_socket.is_linked: + metallic_socket_links = metallic_socket.links[0] + material.node_tree.links.new( + metallic_socket_links.from_socket, + shader_node_socket, + ) + else: + with_emissive_rgb_node = True + orig_metallic_default_value = metallic_socket.default_value + else: + with_emissive_rgb_node = True + + if with_emissive_rgb_node: + rgb_node = material.node_tree.nodes.new(type="ShaderNodeRGB") + rgb_node.outputs[0].default_value = ( + *((orig_metallic_default_value,) * 3), + 1.0, + ) + material.node_tree.links.new( + rgb_node.outputs["Color"], shader_node_socket + ) + ingredients["emissive_rgb_node"] = rgb_node + + return ingredients + + case _: + return {} + + def cleanup( + self, + material: bpy.types.Material, + *, + orig_metallic_from_socket: Optional[bpy.types.NodeSocket] = None, + orig_metallic_default_value: Optional[float] = None, + orig_surface_shader_source: Optional[bpy.types.NodeSocket] = None, + emissive_rgb_node: Optional[bpy.types.Node] = None, + ): + match self: + case self.ALBEDO: + if orig_metallic_default_value or orig_metallic_from_socket: + metallic_socket = ( + [ + node + for node in material.node_tree.nodes + if node.type == "OUTPUT_MATERIAL" + ][0] + .inputs["Surface"] + .links[0] + .inputs["Metallic"] + ) + metallic_socket.default_value = orig_metallic_default_value + if orig_metallic_from_socket: + material.node_tree.links.new( + orig_metallic_from_socket, + metallic_socket, + ) + + case self.METALLIC: + if orig_surface_shader_source: + shader_node_socket = [ + node + for node in material.node_tree.nodes + if node.type == "OUTPUT_MATERIAL" + ][0].inputs["Surface"] + shader_link = shader_node_socket.links[0] + material.node_tree.links.remove(shader_link) + material.node_tree.links.new( + orig_surface_shader_source, + shader_node_socket, + ) + if emissive_rgb_node: + material.node_tree.nodes.remove(emissive_rgb_node) + + class Baker: + """ + Simple wrapper around Blender baking capabilities. + """ + + @staticmethod + def preheat_oven(render_samples: int): + # Only Cycles supports texture baking + bpy.context.scene.render.engine = "CYCLES" + bpy.data.scenes[0].render.engine = "CYCLES" + + # Bake using GPU + bpy.context.preferences.addons["cycles"].preferences.compute_device_type = ( + "CUDA" + ) + bpy.context.preferences.addons["cycles"].preferences.get_devices() + for device in bpy.context.preferences.addons["cycles"].preferences.devices: + device.use = device.type == "CUDA" + bpy.context.scene.render.engine = "CYCLES" + bpy.context.scene.cycles.device = "GPU" + + # Improve performance (only few samples are needed for baking) + bpy.context.scene.cycles.samples = render_samples + + # Consider only the color pass (no environment lighting) + bpy.context.scene.render.bake.use_pass_direct = False + bpy.context.scene.render.bake.use_pass_indirect = False + bpy.context.scene.render.bake.use_pass_color = True + + @classmethod + def bake_into_pbr_material( + cls, + obj: bpy.types.Object, + *, + texture_resolution: int, + ): + # Adjust the recipe according to the object + bpy.context.scene.render.bake.margin = texture_resolution // 64 + + # Unwrap the object if necessary + if not obj.data.uv_layers: + cls._unwrap_uv_on_active_obj(obj, texture_resolution) + + # Get the material + material = obj.data.materials[0] + + # Bake all textures from the recipe + baked_textures = {} + for recipe in ProceduralGenerator.Recipe.enabled_recipes(): + ingredients = recipe.prep(material=material) + + # Create the image node into which the texture will be baked + image_node = cls._create_image_node( + material=material, + recipe=recipe, + texture_resolution=texture_resolution, + ) + + # Bake the texture + bpy.ops.object.bake(type=recipe.bake_type) + baked_textures[recipe] = image_node.image + + # Remove the image node and cleanup the material + material.node_tree.nodes.remove(image_node) + recipe.cleanup(material=material, **ingredients) + + # Create a new PBR material with the baked textures + pbr_material = cls._bake_into_pbr_material( + name=f"PBR_{obj.name}", baked_textures=baked_textures + ) + + # Replace the original material with the new PBR material + obj.data.materials.clear() + obj.data.materials.append(pbr_material) + + @classmethod + def _bake_into_pbr_material( + cls, + name: str, + baked_textures: Dict[ProceduralGenerator.Recipe, bpy.types.Image], + ) -> bpy.types.Material: + # Create a new material + pbr_material = bpy.data.materials.new(name=name) + pbr_material.use_nodes = True + + # Get handles to the nodes and links + nodes = pbr_material.node_tree.nodes + links = pbr_material.node_tree.links + + # Clear the existing nodes + nodes.clear() + + # Create Material Output and Principled BSDF nodes + shader_node = nodes.new(type="ShaderNodeBsdfPrincipled") + shader_node.location = (-300, 0) + output_node = nodes.new(type="ShaderNodeOutputMaterial") + output_node.location = (0, 0) + links.new(shader_node.outputs["BSDF"], output_node.inputs["Surface"]) + + # Create Texture Coordinate and Mapping nodes + texcoord_node = nodes.new(type="ShaderNodeTexCoord") + texcoord_node.location = (-1200, 0) + mapping_node = nodes.new(type="ShaderNodeMapping") + mapping_node.location = (-1000, 0) + links.new(mapping_node.inputs["Vector"], texcoord_node.outputs["UV"]) + + # Create baked textures + for i, (recipe, texture) in enumerate(baked_textures.items()): + # Create Image Texture node + img_texture = nodes.new(type="ShaderNodeTexImage") + img_texture.image = texture + img_texture.location = ( + -800, + 375 * (0.5 * len(baked_textures) + 0.5 - i), + ) + links.new(img_texture.inputs["Vector"], mapping_node.outputs["Vector"]) + + match recipe: + # Normal map requires the Normal Map node + case ProceduralGenerator.Recipe.NORMAL: + normal_map_node = nodes.new(type="ShaderNodeNormalMap") + normal_map_node.location = (-500, -127) + links.new( + img_texture.outputs["Color"], + normal_map_node.inputs["Color"], + ) + links.new( + normal_map_node.outputs["Normal"], + shader_node.inputs[recipe.shader_socket_name], + ) + + # Other textures are directly linked to the shader node + case _: + links.new( + img_texture.outputs["Color"], + shader_node.inputs[recipe.shader_socket_name], + ) + + return pbr_material + + @staticmethod + def _unwrap_uv_on_active_obj(obj: bpy.types.Object, texture_resolution: int): + bpy.ops.object.mode_set(mode="EDIT") + bpy.ops.mesh.select_all(action="SELECT") + + # Try with the default unwrap first, but capture the output in case it fails (printed as warning to stdout) + stdout_str = io.StringIO() + with contextlib.redirect_stdout(stdout_str): + bpy.ops.uv.unwrap() + + # If the default unwrap failed, try with the Smart UV Project + if "Unwrap failed" in stdout_str.getvalue(): + bpy.ops.uv.smart_project( + rotate_method="AXIS_ALIGNED", island_margin=0.5 / texture_resolution + ) + + bpy.ops.object.mode_set(mode="OBJECT") + + @classmethod + def _create_image_node( + cls, + material: bpy.types.Material, + recipe: ProceduralGenerator.Recipe, + texture_resolution: int, + ): + node = material.node_tree.nodes.new("ShaderNodeTexImage") + node.image = cls._create_image_texture( + name=recipe.name.lower(), + texture_resolution=texture_resolution, + color_space=recipe.color_space, + ) + + material.node_tree.nodes.active = node + + return node + + @staticmethod + def _create_image_texture( + name: str, + *, + texture_resolution: int, + color_space: Literal["Non-Color", "sRGB"], + ): + image = bpy.data.images.new( + name, + width=texture_resolution, + height=texture_resolution, + alpha=False, + float_buffer=False, + tiled=False, + ) + image.alpha_mode = "NONE" + image.colorspace_settings.name = color_space + + return image + + class Exporter: + """ + Simple wrapper around Blender export capabilities. + """ + + @staticmethod + def prepare_renderer(): + bpy.ops.object.light_add( + type="SUN", + rotation=(-0.25 * math.pi, -0.4 * math.pi, 0.0), + ) + bpy.context.object.data.energy = 100.0 + bpy.context.object.data.angle = 0.0 + + bpy.ops.object.camera_add(location=(0.0, 0.0, 1.0)) + bpy.context.scene.camera = bpy.context.object + bpy.context.scene.camera.data.type = "ORTHO" + bpy.context.scene.camera.data.clip_start = 1e-06 + + bpy.context.scene.render.film_transparent = True + + @classmethod + def export( + cls, + filepath: Union[str, Path], + *, + ext: str, + render_thumbnail: bool, + thumbnail_resolution: int, + **kwargs, + ): + # Dispatch the export based on the file extension + match ext: + case ".usdz" | ".usdc" | ".usda" | ".usd": + cls.usd_export(filepath=filepath, ext=ext, **kwargs) + case ".sdf": + cls.sdf_export( + filepath=filepath, + render_thumbnail=render_thumbnail, + thumbnail_resolution=thumbnail_resolution, + **kwargs, + ) + case ".abc": + cls.abc_export(filepath=filepath, **kwargs) + case ".fbx": + cls.fbx_export(filepath=filepath, **kwargs) + case ".glb" | ".gltf": + cls.gltf_export(filepath=filepath, ext=ext, **kwargs) + case ".obj": + cls.obj_export(filepath=filepath, **kwargs) + case ".ply": + cls.ply_export(filepath=filepath, **kwargs) + case ".stl": + cls.stl_export(filepath=filepath, **kwargs) + case _: + raise ValueError(f"Unsupported export format: '{ext}'") + + @staticmethod + def abc_export( + filepath: Union[str, Path], + **kwargs, + ): + """ + Export via `bpy.ops.wm.alembic_export()` + """ + + DEFAULT_EXPORT_OVERRIDES = { + "check_existing": False, + "selected": True, + } + + # Prepare the output path + if not isinstance(filepath, Path): + filepath = Path(filepath) + filepath = filepath.with_suffix(".abc").absolute() + + # Forward export kwargs + export_kwargs = deepcopy(DEFAULT_EXPORT_OVERRIDES) + export_kwargs.update(kwargs) + + # Export the Alembic file + bpy.ops.wm.alembic_export( + filepath=filepath.as_posix(), + **export_kwargs, + ) + + @staticmethod + def fbx_export( + filepath: Union[str, Path], + **kwargs, + ): + """ + Export via `bpy.ops.export_scene.fbx()` + """ + + DEFAULT_EXPORT_OVERRIDES = { + "check_existing": False, + "use_selection": True, + } + + # Prepare the output path + if not isinstance(filepath, Path): + filepath = Path(filepath) + filepath = filepath.with_suffix(".fbx").absolute() + + # Forward export kwargs + export_kwargs = deepcopy(DEFAULT_EXPORT_OVERRIDES) + export_kwargs.update(kwargs) + + # Export the FBX file + bpy.ops.export_scene.fbx( + filepath=filepath.as_posix(), + **export_kwargs, + ) + + @staticmethod + def gltf_export( + filepath: Union[str, Path], + *, + ext: str, + **kwargs, + ): + """ + Export via `bpy.ops.export_scene.gltf()` + """ + + DEFAULT_EXPORT_OVERRIDES = { + "check_existing": False, + "use_selection": True, + } + + # Prepare the output path + if not isinstance(filepath, Path): + filepath = Path(filepath) + filepath = filepath.with_suffix(ext).absolute() + + # Forward export kwargs + export_kwargs = deepcopy(DEFAULT_EXPORT_OVERRIDES) + export_kwargs.update(kwargs) + + # Export the GLTF file + bpy.ops.wm.usd_export( + filepath=filepath.as_posix(), + **export_kwargs, + ) + + @staticmethod + def obj_export( + filepath: Union[str, Path], + **kwargs, + ): + """ + Export via `bpy.ops.wm.obj_export()` + """ + + DEFAULT_EXPORT_OVERRIDES = { + "check_existing": False, + "export_selected_objects": True, + } + + # Prepare the output path + if not isinstance(filepath, Path): + filepath = Path(filepath) + filepath = filepath.with_suffix(".obj").absolute() + + # Forward export kwargs + export_kwargs = deepcopy(DEFAULT_EXPORT_OVERRIDES) + export_kwargs.update(kwargs) + + # Export the Wavefront OBJ file + bpy.ops.wm.obj_export( + filepath=filepath.as_posix(), + **export_kwargs, + ) + + @staticmethod + def ply_export( + filepath: Union[str, Path], + **kwargs, + ): + """ + Export via `bpy.ops.wm.ply_export()` + """ + + DEFAULT_EXPORT_OVERRIDES = { + "check_existing": False, + "export_selected_objects": True, + } + + # Prepare the output path + if not isinstance(filepath, Path): + filepath = Path(filepath) + filepath = filepath.with_suffix(".ply").absolute() + + # Forward export kwargs + export_kwargs = deepcopy(DEFAULT_EXPORT_OVERRIDES) + export_kwargs.update(kwargs) + + # Export the PLY file + bpy.ops.wm.ply_export( + filepath=filepath.as_posix(), + **export_kwargs, + ) + + @staticmethod + def sdf_export( + filepath: Union[str, Path], + *, + render_thumbnail: bool, + thumbnail_resolution: int, + **kwargs, + ): + """ + Export an SDF model with a corresponding GLTF mesh. + """ + + DEFAULT_EXPORT_OVERRIDES = { + "check_existing": False, + "use_selection": True, + } + + # Prepare the output path + if not isinstance(filepath, Path): + filepath = Path(filepath) + filepath = filepath.absolute() + model_name = filepath.stem + filepath_model = Path(path.join(filepath, "model.sdf")) + filepath_config = Path(path.join(filepath, "model.config")) + filepath_mesh = Path(path.join(filepath, "meshes", f"{model_name}.glb")) + filepath_thumbnail = Path(path.join(filepath, "thumbnails", "0.png")) + + # Create parent directories + os.makedirs(name=filepath, exist_ok=True) + os.makedirs(name=filepath_mesh.parent, exist_ok=True) + os.makedirs(name=filepath_thumbnail.parent, exist_ok=True) + + # Forward export kwargs + export_kwargs = deepcopy(DEFAULT_EXPORT_OVERRIDES) + export_kwargs.update(kwargs) + + # Export the USD file + bpy.ops.export_scene.gltf( + filepath=filepath_mesh.as_posix(), + **export_kwargs, + ) + + # Write the SDF model file + sdf = xml_et.Element("sdf", attrib={"version": "1.9"}) + model = xml_et.SubElement(sdf, "model", attrib={"name": model_name}) + link = xml_et.SubElement( + model, "link", attrib={"name": f"{model_name}_link"} + ) + pose = xml_et.SubElement(link, "pose") + pose.text = f"0 0 0 {math.pi/2} 0 0" + visual = xml_et.SubElement( + link, "visual", attrib={"name": f"{model_name}_visual"} + ) + visual_geometry = xml_et.SubElement(visual, "geometry") + visual_mesh = xml_et.SubElement(visual_geometry, "mesh") + visual_mesh_uri = xml_et.SubElement(visual_mesh, "uri") + visual_mesh_uri.text = path.relpath(start=filepath, path=filepath_mesh) + collision = xml_et.SubElement( + link, "collision", attrib={"name": f"{model_name}_collision"} + ) + collision_geometry = xml_et.SubElement(collision, "geometry") + collision_mesh = xml_et.SubElement(collision_geometry, "mesh") + collision_mesh_uri = xml_et.SubElement(collision_mesh, "uri") + collision_mesh_uri.text = path.relpath(start=filepath, path=filepath_mesh) + xml_str = xml_dom.parseString(xml_et.tostring(sdf, encoding="unicode")) + xml_str = xml_str.toprettyxml(indent=" ") + sdf_file = open(filepath_model.as_posix(), "w") + sdf_file.write(xml_str) + sdf_file.close() + + # Write the model config file + cfg = xml_et.Element("model") + name = xml_et.SubElement(cfg, "name") + name.text = model_name + version = xml_et.SubElement(cfg, "version") + version.text = "1" + sdf_tag = xml_et.SubElement(cfg, "sdf", attrib={"version": "1.9"}) + sdf_tag.text = "model.sdf" + author = xml_et.SubElement(cfg, "author") + producer = xml_et.SubElement(author, "producer") + producer.text = f"Blender {bpy.app.version_string}" + description = xml_et.SubElement(cfg, "description") + description.text = "Procedurally generated model" + xml_str = xml_dom.parseString(xml_et.tostring(cfg, encoding="unicode")) + xml_str = xml_str.toprettyxml(indent=" ") + model_cfg_file = open(filepath_config.as_posix(), "w") + model_cfg_file.write(xml_str) + model_cfg_file.close() + + # Render the thumbnail + if render_thumbnail: + obj_bbox_center = bpy.context.object.matrix_world @ ( + sum( + (mathutils.Vector(b) for b in bpy.context.object.bound_box), + mathutils.Vector(), + ) + / 8.0 + ) + obj_dims = bpy.context.object.dimensions + bpy.context.scene.camera.location = ( + obj_bbox_center[0], + obj_bbox_center[1], + obj_bbox_center[2] + obj_dims[2], + ) + max_xy_dim = max(obj_dims[:2]) + bpy.context.scene.camera.data.ortho_scale = max_xy_dim + bpy.context.scene.camera.data.clip_end = 2.0 * obj_dims[2] + bpy.context.scene.render.resolution_x = math.ceil( + thumbnail_resolution * max_xy_dim / obj_dims[1] + ) + bpy.context.scene.render.resolution_y = math.ceil( + thumbnail_resolution * max_xy_dim / obj_dims[0] + ) + bpy.context.scene.render.filepath = filepath_thumbnail.as_posix() + bpy.ops.render.render(write_still=True) + + @staticmethod + def stl_export( + filepath: Union[str, Path], + **kwargs, + ): + """ + Export via `bpy.ops.wm.stl_export()` + """ + + DEFAULT_EXPORT_OVERRIDES = { + "check_existing": False, + "export_selected_objects": True, + } + + # Prepare the output path + if not isinstance(filepath, Path): + filepath = Path(filepath) + filepath = filepath.with_suffix(".stl").absolute() + + # Forward export kwargs + export_kwargs = deepcopy(DEFAULT_EXPORT_OVERRIDES) + export_kwargs.update(kwargs) + + # Export the STL file + bpy.ops.wm.stl_export( + filepath=filepath.as_posix(), + **export_kwargs, + ) + + @staticmethod + def usd_export( + filepath: Union[str, Path], + *, + ext: str, + **kwargs, + ): + """ + Export via `bpy.ops.wm.usd_export()` + """ + + DEFAULT_EXPORT_OVERRIDES = { + "check_existing": False, + "selected_objects_only": True, + "author_blender_name": False, + "use_instancing": True, + } + + # Prepare the output path + if not isinstance(filepath, Path): + filepath = Path(filepath) + filepath = filepath.with_suffix(ext).absolute() + + # Forward export kwargs + export_kwargs = deepcopy(DEFAULT_EXPORT_OVERRIDES) + export_kwargs.update(kwargs) + + # Export the USD file + bpy.ops.wm.usd_export( + filepath=filepath.as_posix(), + **export_kwargs, + ) + + +### String utils ### +REGEX_CANONICALIZE_STR_PATTERN: re.Pattern = re.compile("[\W_]+") + + +def canonicalize_str(input: str) -> str: + """ + Canonicalizes a string by converting it to lowercase and removing unwanted characters. + + This function processes the input string to ensure it is in a standardized format, making it suitable for consistent usage in applications. It utilizes a predefined regular expression pattern to eliminate any characters that do not meet the specified criteria. + + Args: + input (str): The string to be canonicalized. + + Returns: + str: The canonicalized version of the input string. + """ + return REGEX_CANONICALIZE_STR_PATTERN.sub("", input.lower()) + + +### Misc utils ### +def print_bpy(msg: Any, file: Optional[TextIO] = sys.stdout, *args, **kwargs): + """ + Helper print function that also provides output inside the Blender console in addition to the system console. + """ + + print(msg, file=file, *args, **kwargs) + for window in bpy.context.window_manager.windows: + for area in window.screen.areas: + if area.type == "CONSOLE": + with bpy.context.temp_override( + window=window, screen=window.screen, area=area + ): + bpy.ops.console.scrollback_append( + text=str(msg), + type="ERROR" if file == sys.stderr else "OUTPUT", + ) + + +def verify_requirements(): + VERSION_BPY_MIN: Tuple[int, int] = (4, 2) + if ( + bpy.app.version[0] != VERSION_BPY_MIN[0] + or bpy.app.version[1] < VERSION_BPY_MIN[1] + ): + print_bpy( + f"[WARNING]: Blender {bpy.app.version_string} is likely incompatible with this script (written for Blender {VERSION_BPY_MIN[0]}.{VERSION_BPY_MIN[1]})", + file=sys.stderr, + ) + + +### CLI ### +def parse_cli_args() -> argparse.Namespace: + """ + Parse command-line arguments for this script. + """ + + parser = argparse.ArgumentParser( + description="Generate procedural dataset of USD assets using Blender", + usage=f"{sys.argv[0] if sys.argv[0].endswith('blender') else 'blender'} --python {path.realpath(__file__)} -- [options]", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + argument_default=argparse.SUPPRESS, + ) + + group = parser.add_argument_group("Input") + group.add_argument( + "-i", + "--autorun_scripts", + type=str, + nargs="*", + help="List of Blender scripts to execute", + required=True, + ) + + group = parser.add_argument_group("Output") + group.add_argument( + "-o", + "--outdir", + type=str, + help="The output directory", + required=True, + ) + group.add_argument( + "--name", + type=str, + help="The base name of the exported models", + default="model", + ) + group.add_argument( + "--ext", + type=str, + help="""The file extension of the exported models (ABC: '.abc',\ + FBX: '.fbx',\ + GLTF: '.glb|.gltf',\ + OBJ: '.obj',\ + PLY: '.ply',\ + SDF: '.sdf',\ + STL: '.stl',\ + USD: '.usd|.usda|.usdc|.usdz')""", + default=".usdz", + ) + group.add_argument( + "--overwrite_min_age", + type=int, + help="Number of seconds after which to overwrite the generated assets if they already exist (disabled if negative)", + default=0, + ) + + group = parser.add_argument_group("Generator") + group.add_argument( + "-s", + "--seed", + type=int, + help="The initial seed of the random number generator", + default=0, + ) + group.add_argument( + "-n", + "--num_assets", + type=int, + help="Number of assets to generate", + default=1, + ) + + group = parser.add_argument_group("Export") + group.add_argument( + "--export_kwargs", + type=json.loads, + help="Keyword arguments for the USD export", + default={}, + ) + group.add_argument( + "--render_thumbnail", + action="store_true", + help="If specified, generate a thumbnail image for each exported model (only for SDF export)", + default=False, + ) + group.add_argument( + "--thumbnail_resolution", + type=int, + help="Resolution of the generated thumbnail images", + default=512, + ) + + group = parser.add_argument_group("Geometry") + group.add_argument( + "--geometry_nodes", + type=json.loads, + help="List of Geometry Nodes modifiers from `--autorun_scripts` for generating the geometry, with an optional dictionary for configuring their inputs", + required=True, + ) + group.add_argument( + "--decimate_angle_limit", + type=float, + help="If specified, decimate the generated geometry to the specified target angle limit", + default=None, + ) + group.add_argument( + "--decimate_face_count", + type=int, + help="If specified, decimate the generated geometry to the specified target face count", + default=None, + ) + + group = parser.add_argument_group("Material") + group.add_argument( + "--material", + type=str, + help="Material of the generated models from `--autorun_scripts`, which will be baked as PBR textures into the USD file", + default=None, + ) + group.add_argument( + "--texture_resolution", + type=int, + help="Resolution of the baked PBR textures", + default=1024, + ) + group.add_argument( + "--render_samples", + type=int, + help="Number of samples to use for rendering the PBR textures", + default=4, + ) + + if "--" in sys.argv: + args = parser.parse_args(sys.argv[sys.argv.index("--") + 1 :]) + else: + args, unknown_args = parser.parse_known_args() + if unknown_args: + print_bpy( + f"[WARNING]: Unknown args: {unknown_args}", + file=sys.stderr, + ) + print_bpy( + '[HINT]: Consider delimiting your args for Python script with "--"' + ) + + return args + + +if __name__ == "__main__": + main(**vars(parse_cli_args()))