a fork of shap-e for gc
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

685 lines
26 KiB

"""
Script to run within blender.
Provide arguments after `--`.
For example: `blender -b -P blender_script.py -- --help`
"""
import argparse
import json
import math
import os
import random
import sys
import bpy
from mathutils import Vector
from mathutils.noise import random_unit_vector
MAX_DEPTH = 5.0
FORMAT_VERSION = 6
# Set by main(), these constants are passed to the script to avoid
# duplicating them across multiple files.
UNIFORM_LIGHT_DIRECTION = None
BASIC_AMBIENT_COLOR = None
BASIC_DIFFUSE_COLOR = None
def clear_scene():
bpy.ops.object.select_all(action="SELECT")
bpy.ops.object.delete()
def clear_lights():
bpy.ops.object.select_all(action="DESELECT")
for obj in bpy.context.scene.objects.values():
if isinstance(obj.data, bpy.types.Light):
obj.select_set(True)
bpy.ops.object.delete()
def import_model(path):
clear_scene()
_, ext = os.path.splitext(path)
ext = ext.lower()
if ext == ".obj":
bpy.ops.import_scene.obj(filepath=path)
elif ext in [".glb", ".gltf"]:
bpy.ops.import_scene.gltf(filepath=path)
elif ext == ".stl":
bpy.ops.import_mesh.stl(filepath=path)
elif ext == ".fbx":
bpy.ops.import_scene.fbx(filepath=path)
elif ext == ".dae":
bpy.ops.wm.collada_import(filepath=path)
elif ext == ".ply":
bpy.ops.import_mesh.ply(filepath=path)
else:
raise RuntimeError(f"unexpected extension: {ext}")
def scene_root_objects():
for obj in bpy.context.scene.objects.values():
if not obj.parent:
yield obj
def scene_bbox(single_obj=None, ignore_matrix=False):
bbox_min = (math.inf,) * 3
bbox_max = (-math.inf,) * 3
found = False
for obj in scene_meshes() if single_obj is None else [single_obj]:
found = True
for coord in obj.bound_box:
coord = Vector(coord)
if not ignore_matrix:
coord = obj.matrix_world @ coord
bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
if not found:
raise RuntimeError("no objects in scene to compute bounding box for")
return Vector(bbox_min), Vector(bbox_max)
def scene_meshes():
for obj in bpy.context.scene.objects.values():
if isinstance(obj.data, (bpy.types.Mesh)):
yield obj
def normalize_scene():
if len(list(scene_root_objects())) > 1:
# Create an empty object to be used as a parent for all root objects
parent_empty = bpy.data.objects.new("ParentEmpty", None)
bpy.context.scene.collection.objects.link(parent_empty)
# Parent all root objects to the empty object
for obj in scene_root_objects():
if obj != parent_empty:
obj.parent = parent_empty
bbox_min, bbox_max = scene_bbox()
scale = 1 / max(bbox_max - bbox_min)
for obj in scene_root_objects():
obj.scale = obj.scale * scale
# Apply scale to matrix_world.
bpy.context.view_layer.update()
bbox_min, bbox_max = scene_bbox()
offset = -(bbox_min + bbox_max) / 2
for obj in scene_root_objects():
obj.matrix_world.translation += offset
bpy.ops.object.select_all(action="DESELECT")
def create_camera():
# https://b3d.interplanety.org/en/how-to-create-camera-through-the-blender-python-api/
camera_data = bpy.data.cameras.new(name="Camera")
camera_object = bpy.data.objects.new("Camera", camera_data)
bpy.context.scene.collection.objects.link(camera_object)
bpy.context.scene.camera = camera_object
def set_camera(direction, camera_dist=2.0):
camera_pos = -camera_dist * direction
bpy.context.scene.camera.location = camera_pos
# https://blender.stackexchange.com/questions/5210/pointing-the-camera-in-a-particular-direction-programmatically
rot_quat = direction.to_track_quat("-Z", "Y")
bpy.context.scene.camera.rotation_euler = rot_quat.to_euler()
bpy.context.view_layer.update()
def randomize_camera(camera_dist=2.0):
direction = random_unit_vector()
set_camera(direction, camera_dist=camera_dist)
def pan_camera(time, axis="Z", camera_dist=2.0, elevation=0.1):
angle = time * math.pi * 2
direction = [-math.cos(angle), -math.sin(angle), elevation]
assert axis in ["X", "Y", "Z"]
if axis == "X":
direction = [direction[2], *direction[:2]]
elif axis == "Y":
direction = [direction[0], elevation, direction[1]]
direction = Vector(direction).normalized()
set_camera(direction, camera_dist=camera_dist)
def place_camera(time, camera_pose_mode="random", camera_dist_min=2.0, camera_dist_max=2.0):
camera_dist = random.uniform(camera_dist_min, camera_dist_max)
if camera_pose_mode == "random":
randomize_camera(camera_dist=camera_dist)
elif camera_pose_mode == "z-circular":
pan_camera(time, axis="Z", camera_dist=camera_dist)
elif camera_pose_mode == "z-circular-elevated":
pan_camera(time, axis="Z", camera_dist=camera_dist, elevation=-0.2617993878)
else:
raise ValueError(f"Unknown camera pose mode: {camera_pose_mode}")
def create_light(location, energy=1.0, angle=0.5 * math.pi / 180):
# https://blender.stackexchange.com/questions/215624/how-to-create-a-light-with-the-python-api-in-blender-2-92
light_data = bpy.data.lights.new(name="Light", type="SUN")
light_data.energy = energy
light_data.angle = angle
light_object = bpy.data.objects.new(name="Light", object_data=light_data)
direction = -location
rot_quat = direction.to_track_quat("-Z", "Y")
light_object.rotation_euler = rot_quat.to_euler()
bpy.context.view_layer.update()
bpy.context.collection.objects.link(light_object)
light_object.location = location
def create_random_lights(count=4, distance=2.0, energy=1.5):
clear_lights()
for _ in range(count):
create_light(random_unit_vector() * distance, energy=energy)
def create_camera_light():
clear_lights()
create_light(bpy.context.scene.camera.location, energy=5.0)
def create_uniform_light(backend):
clear_lights()
# Random direction to decorrelate axis-aligned sides.
pos = Vector(UNIFORM_LIGHT_DIRECTION)
angle = 0.0092 if backend == "CYCLES" else math.pi
create_light(pos, energy=5.0, angle=angle)
create_light(-pos, energy=5.0, angle=angle)
def create_vertex_color_shaders():
# By default, Blender will ignore vertex colors in both the
# Eevee and Cycles backends, since these colors aren't
# associated with a material.
#
# What we do here is create a simple material shader and link
# the vertex color to the material color.
for obj in bpy.context.scene.objects.values():
if not isinstance(obj.data, (bpy.types.Mesh)):
continue
if len(obj.data.materials):
# We don't want to override any existing materials.
continue
color_keys = (obj.data.vertex_colors or {}).keys()
if not len(color_keys):
# Many objects will have no materials *or* vertex colors.
continue
mat = bpy.data.materials.new(name="VertexColored")
mat.use_nodes = True
# There should be a Principled BSDF by default.
bsdf_node = None
for node in mat.node_tree.nodes:
if node.type == "BSDF_PRINCIPLED":
bsdf_node = node
assert bsdf_node is not None, "material has no Principled BSDF node to modify"
socket_map = {}
for input in bsdf_node.inputs:
socket_map[input.name] = input
# Make sure nothing lights the object except for the diffuse color.
socket_map["Specular"].default_value = 0.0
socket_map["Roughness"].default_value = 1.0
v_color = mat.node_tree.nodes.new("ShaderNodeVertexColor")
v_color.layer_name = color_keys[0]
mat.node_tree.links.new(v_color.outputs[0], socket_map["Base Color"])
obj.data.materials.append(mat)
def create_default_materials():
for obj in bpy.context.scene.objects.values():
if isinstance(obj.data, (bpy.types.Mesh)):
if not len(obj.data.materials):
mat = bpy.data.materials.new(name="DefaultMaterial")
mat.use_nodes = True
obj.data.materials.append(mat)
def find_materials():
all_materials = set()
for obj in bpy.context.scene.objects.values():
if not isinstance(obj.data, bpy.types.Mesh):
continue
for mat in obj.data.materials:
all_materials.add(mat)
return all_materials
def delete_all_materials():
for obj in bpy.context.scene.objects.values():
if isinstance(obj.data, bpy.types.Mesh):
# https://blender.stackexchange.com/questions/146714/removing-all-material-slots-in-one-go
obj.data.materials.clear()
def setup_material_extraction_shaders(capturing_material_alpha: bool):
"""
Change every material to emit texture colors (or alpha) rather than having
an actual reflective color. Returns a function to undo the changes to the
materials.
"""
# Objects can share materials, so we first find all of the
# materials in the project, and then modify them each once.
undo_fns = []
for mat in find_materials():
undo_fn = setup_material_extraction_shader_for_material(mat, capturing_material_alpha)
if undo_fn is not None:
undo_fns.append(undo_fn)
return lambda: [undo_fn() for undo_fn in undo_fns]
def setup_material_extraction_shader_for_material(mat, capturing_material_alpha: bool):
mat.use_nodes = True
# By default, most imported models should use the regular
# "Principled BSDF" material, so we should always find this.
# If not, this shader manipulation logic won't work.
bsdf_node = None
for node in mat.node_tree.nodes:
if node.type == "BSDF_PRINCIPLED":
bsdf_node = node
assert bsdf_node is not None, "material has no Principled BSDF node to modify"
socket_map = {}
for input in bsdf_node.inputs:
socket_map[input.name] = input
for name in ["Base Color", "Emission", "Emission Strength", "Alpha", "Specular"]:
assert name in socket_map.keys(), f"{name} not in {list(socket_map.keys())}"
old_base_color = get_socket_value(mat.node_tree, socket_map["Base Color"])
old_alpha = get_socket_value(mat.node_tree, socket_map["Alpha"])
old_emission = get_socket_value(mat.node_tree, socket_map["Emission"])
old_emission_strength = get_socket_value(mat.node_tree, socket_map["Emission Strength"])
old_specular = get_socket_value(mat.node_tree, socket_map["Specular"])
# Make sure the base color of all objects is black and the opacity
# is 1, so that we are effectively just telling the shader what color
# to make the pixels.
clear_socket_input(mat.node_tree, socket_map["Base Color"])
socket_map["Base Color"].default_value = [0, 0, 0, 1]
clear_socket_input(mat.node_tree, socket_map["Alpha"])
socket_map["Alpha"].default_value = 1
clear_socket_input(mat.node_tree, socket_map["Specular"])
socket_map["Specular"].default_value = 0.0
old_blend_method = mat.blend_method
mat.blend_method = "OPAQUE"
if capturing_material_alpha:
set_socket_value(mat.node_tree, socket_map["Emission"], old_alpha)
else:
set_socket_value(mat.node_tree, socket_map["Emission"], old_base_color)
clear_socket_input(mat.node_tree, socket_map["Emission Strength"])
socket_map["Emission Strength"].default_value = 1.0
def undo_fn():
mat.blend_method = old_blend_method
set_socket_value(mat.node_tree, socket_map["Base Color"], old_base_color)
set_socket_value(mat.node_tree, socket_map["Alpha"], old_alpha)
set_socket_value(mat.node_tree, socket_map["Emission"], old_emission)
set_socket_value(mat.node_tree, socket_map["Emission Strength"], old_emission_strength)
set_socket_value(mat.node_tree, socket_map["Specular"], old_specular)
return undo_fn
def get_socket_value(tree, socket):
default = socket.default_value
if not isinstance(default, float):
default = list(default)
for link in tree.links:
if link.to_socket == socket:
return (link.from_socket, default)
return (None, default)
def clear_socket_input(tree, socket):
for link in list(tree.links):
if link.to_socket == socket:
tree.links.remove(link)
def set_socket_value(tree, socket, socket_and_default):
clear_socket_input(tree, socket)
old_source_socket, default = socket_and_default
if isinstance(default, float) and not isinstance(socket.default_value, float):
# Codepath for setting Emission to a previous alpha value.
socket.default_value = [default] * 3 + [1.0]
else:
socket.default_value = default
if old_source_socket is not None:
tree.links.new(old_source_socket, socket)
def setup_nodes(output_path, capturing_material_alpha: bool = False, basic_lighting: bool = False):
tree = bpy.context.scene.node_tree
links = tree.links
for node in tree.nodes:
tree.nodes.remove(node)
# Helpers to perform math on links and constants.
def node_op(op: str, *args, clamp=False):
node = tree.nodes.new(type="CompositorNodeMath")
node.operation = op
if clamp:
node.use_clamp = True
for i, arg in enumerate(args):
if isinstance(arg, (int, float)):
node.inputs[i].default_value = arg
else:
links.new(arg, node.inputs[i])
return node.outputs[0]
def node_clamp(x, maximum=1.0):
return node_op("MINIMUM", x, maximum)
def node_mul(x, y, **kwargs):
return node_op("MULTIPLY", x, y, **kwargs)
def node_add(x, y, **kwargs):
return node_op("ADD", x, y, **kwargs)
def node_abs(x, **kwargs):
return node_op("ABSOLUTE", x, **kwargs)
input_node = tree.nodes.new(type="CompositorNodeRLayers")
input_node.scene = bpy.context.scene
input_sockets = {}
for output in input_node.outputs:
input_sockets[output.name] = output
if capturing_material_alpha:
color_socket = input_sockets["Image"]
else:
raw_color_socket = input_sockets["Image"]
if basic_lighting:
# Compute diffuse lighting
normal_xyz = tree.nodes.new(type="CompositorNodeSeparateXYZ")
tree.links.new(input_sockets["Normal"], normal_xyz.inputs[0])
normal_x, normal_y, normal_z = [normal_xyz.outputs[i] for i in range(3)]
dot = node_add(
node_mul(UNIFORM_LIGHT_DIRECTION[0], normal_x),
node_add(
node_mul(UNIFORM_LIGHT_DIRECTION[1], normal_y),
node_mul(UNIFORM_LIGHT_DIRECTION[2], normal_z),
),
)
diffuse = node_abs(dot)
# Compute ambient + diffuse lighting
brightness = node_add(BASIC_AMBIENT_COLOR, node_mul(BASIC_DIFFUSE_COLOR, diffuse))
# Modulate the RGB channels using the total brightness.
rgba_node = tree.nodes.new(type="CompositorNodeSepRGBA")
tree.links.new(raw_color_socket, rgba_node.inputs[0])
combine_node = tree.nodes.new(type="CompositorNodeCombRGBA")
for i in range(3):
tree.links.new(node_mul(rgba_node.outputs[i], brightness), combine_node.inputs[i])
tree.links.new(rgba_node.outputs[3], combine_node.inputs[3])
raw_color_socket = combine_node.outputs[0]
# We apply sRGB here so that our fixed-point depth map and material
# alpha values are not sRGB, and so that we perform ambient+diffuse
# lighting in linear RGB space.
color_node = tree.nodes.new(type="CompositorNodeConvertColorSpace")
color_node.from_color_space = "Linear"
color_node.to_color_space = "sRGB"
tree.links.new(raw_color_socket, color_node.inputs[0])
color_socket = color_node.outputs[0]
split_node = tree.nodes.new(type="CompositorNodeSepRGBA")
tree.links.new(color_socket, split_node.inputs[0])
# Create separate file output nodes for every channel we care about.
# The process calling this script must decide how to recombine these
# channels, possibly into a single image.
for i, channel in enumerate("rgba") if not capturing_material_alpha else [(0, "MatAlpha")]:
output_node = tree.nodes.new(type="CompositorNodeOutputFile")
output_node.base_path = f"{output_path}_{channel}"
links.new(split_node.outputs[i], output_node.inputs[0])
if capturing_material_alpha:
# No need to re-write depth here.
return
depth_out = node_clamp(node_mul(input_sockets["Depth"], 1 / MAX_DEPTH))
output_node = tree.nodes.new(type="CompositorNodeOutputFile")
output_node.base_path = f"{output_path}_depth"
links.new(depth_out, output_node.inputs[0])
def render_scene(output_path, fast_mode: bool, extract_material: bool, basic_lighting: bool):
use_workbench = bpy.context.scene.render.engine == "BLENDER_WORKBENCH"
if use_workbench:
# We must use a different engine to compute depth maps.
bpy.context.scene.render.engine = "BLENDER_EEVEE"
bpy.context.scene.eevee.taa_render_samples = 1 # faster, since we discard image.
if fast_mode:
if bpy.context.scene.render.engine == "BLENDER_EEVEE":
bpy.context.scene.eevee.taa_render_samples = 1
elif bpy.context.scene.render.engine == "CYCLES":
bpy.context.scene.cycles.samples = 256
else:
if bpy.context.scene.render.engine == "CYCLES":
# We should still impose a per-frame time limit
# so that we don't timeout completely.
bpy.context.scene.cycles.time_limit = 40
bpy.context.view_layer.update()
bpy.context.scene.use_nodes = True
bpy.context.scene.view_layers["ViewLayer"].use_pass_z = True
if basic_lighting:
bpy.context.scene.view_layers["ViewLayer"].use_pass_normal = True
bpy.context.scene.view_settings.view_transform = "Raw" # sRGB done in graph nodes
bpy.context.scene.render.film_transparent = True
bpy.context.scene.render.resolution_x = 512
bpy.context.scene.render.resolution_y = 512
bpy.context.scene.render.image_settings.file_format = "PNG"
bpy.context.scene.render.image_settings.color_mode = "BW"
bpy.context.scene.render.image_settings.color_depth = "16"
bpy.context.scene.render.filepath = output_path
if extract_material:
for do_alpha in [False, True]:
undo_fn = setup_material_extraction_shaders(capturing_material_alpha=do_alpha)
setup_nodes(output_path, capturing_material_alpha=do_alpha)
bpy.ops.render.render(write_still=True)
undo_fn()
else:
setup_nodes(output_path, basic_lighting=basic_lighting)
bpy.ops.render.render(write_still=True)
# The output images must be moved from their own sub-directories, or
# discarded if we are using workbench for the color.
for channel_name in ["r", "g", "b", "a", "depth", *(["MatAlpha"] if extract_material else [])]:
sub_dir = f"{output_path}_{channel_name}"
image_path = os.path.join(sub_dir, os.listdir(sub_dir)[0])
name, ext = os.path.splitext(output_path)
if channel_name == "depth" or not use_workbench:
os.rename(image_path, f"{name}_{channel_name}{ext}")
else:
os.remove(image_path)
os.removedirs(sub_dir)
if use_workbench:
# Re-render RGBA using workbench with texture mode, since this seems
# to show the most reasonable colors when lighting is broken.
bpy.context.scene.use_nodes = False
bpy.context.scene.render.engine = "BLENDER_WORKBENCH"
bpy.context.scene.render.image_settings.color_mode = "RGBA"
bpy.context.scene.render.image_settings.color_depth = "8"
bpy.context.scene.display.shading.color_type = "TEXTURE"
bpy.context.scene.display.shading.light = "FLAT"
if fast_mode:
# Single pass anti-aliasing.
bpy.context.scene.display.render_aa = "FXAA"
os.remove(output_path)
bpy.ops.render.render(write_still=True)
bpy.context.scene.render.image_settings.color_mode = "BW"
bpy.context.scene.render.image_settings.color_depth = "16"
def scene_fov():
x_fov = bpy.context.scene.camera.data.angle_x
y_fov = bpy.context.scene.camera.data.angle_y
width = bpy.context.scene.render.resolution_x
height = bpy.context.scene.render.resolution_y
if bpy.context.scene.camera.data.angle == x_fov:
y_fov = 2 * math.atan(math.tan(x_fov / 2) * height / width)
else:
x_fov = 2 * math.atan(math.tan(y_fov / 2) * width / height)
return x_fov, y_fov
def write_camera_metadata(path):
x_fov, y_fov = scene_fov()
bbox_min, bbox_max = scene_bbox()
matrix = bpy.context.scene.camera.matrix_world
with open(path, "w") as f:
json.dump(
dict(
format_version=FORMAT_VERSION,
max_depth=MAX_DEPTH,
bbox=[list(bbox_min), list(bbox_max)],
origin=list(matrix.col[3])[:3],
x_fov=x_fov,
y_fov=y_fov,
x=list(matrix.col[0])[:3],
y=list(-matrix.col[1])[:3],
z=list(-matrix.col[2])[:3],
),
f,
)
def save_rendering_dataset(
input_path: str,
output_path: str,
num_images: int,
backend: str,
light_mode: str,
camera_pose: str,
camera_dist_min: float,
camera_dist_max: float,
fast_mode: bool,
extract_material: bool,
delete_material: bool,
):
assert light_mode in ["random", "uniform", "camera", "basic"]
assert camera_pose in ["random", "z-circular", "z-circular-elevated"]
basic_lighting = light_mode == "basic"
assert not (basic_lighting and extract_material), "cannot extract material with basic lighting"
assert not (delete_material and extract_material), "cannot extract material and delete it"
import_model(input_path)
bpy.context.scene.render.engine = backend
normalize_scene()
if light_mode == "random":
create_random_lights()
elif light_mode == "uniform":
create_uniform_light(backend)
create_camera()
create_vertex_color_shaders()
if delete_material:
delete_all_materials()
if extract_material or basic_lighting:
create_default_materials()
if basic_lighting:
# Make sure materials are uniformly lit, so that we can light
# them in the output shader.
setup_material_extraction_shaders(capturing_material_alpha=False)
for i in range(num_images):
t = i / max(num_images - 1, 1) # same as np.linspace(0, 1, num_images)
place_camera(
t,
camera_pose_mode=camera_pose,
camera_dist_min=camera_dist_min,
camera_dist_max=camera_dist_max,
)
if light_mode == "camera":
create_camera_light()
render_scene(
os.path.join(output_path, f"{i:05}.png"),
fast_mode=fast_mode,
extract_material=extract_material,
basic_lighting=basic_lighting,
)
write_camera_metadata(os.path.join(output_path, f"{i:05}.json"))
with open(os.path.join(output_path, "info.json"), "w") as f:
info = dict(
backend=backend,
light_mode=light_mode,
fast_mode=fast_mode,
extract_material=extract_material,
format_version=FORMAT_VERSION,
channels=["R", "G", "B", "A", "D", *(["MatAlpha"] if extract_material else [])],
scale=0.5, # The scene is bounded by [-scale, scale].
)
json.dump(info, f)
def main():
global UNIFORM_LIGHT_DIRECTION, BASIC_AMBIENT_COLOR, BASIC_DIFFUSE_COLOR
try:
dash_index = sys.argv.index("--")
except ValueError as exc:
raise ValueError("arguments must be preceded by '--'") from exc
raw_args = sys.argv[dash_index + 1 :]
parser = argparse.ArgumentParser()
parser.add_argument("--input_path", required=True, type=str)
parser.add_argument("--output_path", required=True, type=str)
parser.add_argument("--num_images", required=True, type=int)
parser.add_argument("--backend", type=str, default="BLENDER_EEVEE")
parser.add_argument("--light_mode", type=str, default="random")
parser.add_argument("--camera_pose", type=str, default="random")
parser.add_argument("--camera_dist_min", type=float, default=2.0)
parser.add_argument("--camera_dist_max", type=float, default=2.0)
parser.add_argument("--fast_mode", action="store_true")
parser.add_argument("--extract_material", action="store_true")
parser.add_argument("--delete_material", action="store_true")
# Prevent constants from being repeated.
parser.add_argument("--uniform_light_direction", required=True, type=float, nargs="+")
parser.add_argument("--basic_ambient", required=True, type=float)
parser.add_argument("--basic_diffuse", required=True, type=float)
args = parser.parse_args(raw_args)
UNIFORM_LIGHT_DIRECTION = args.uniform_light_direction
BASIC_AMBIENT_COLOR = args.basic_ambient
BASIC_DIFFUSE_COLOR = args.basic_diffuse
save_rendering_dataset(
input_path=args.input_path,
output_path=args.output_path,
num_images=args.num_images,
backend=args.backend,
light_mode=args.light_mode,
camera_pose=args.camera_pose,
camera_dist_min=args.camera_dist_min,
camera_dist_max=args.camera_dist_max,
fast_mode=args.fast_mode,
extract_material=args.extract_material,
delete_material=args.delete_material,
)
main()