scripts/render_views.py

"""Blender CPU review renders, including labeled component previews."""
import argparse
import json
import math
from pathlib import Path
import sys

import bpy
from mathutils import Vector

sys.path.insert(0, str(Path(__file__).resolve().parent))
from render_settings import configure_cpu_render

parser = argparse.ArgumentParser()
parser.add_argument('--input', required=True)
parser.add_argument('--output', required=True)
parser.add_argument('--components', action='store_true')
parser.add_argument('--material-mode', choices=['original', 'clay', 'basecolor'], default='original')
args = parser.parse_args(sys.argv[sys.argv.index('--') + 1:])
folder = Path(args.output)
folder.mkdir(parents=True, exist_ok=True)
bpy.ops.object.select_all(action='SELECT')
bpy.ops.object.delete(use_global=False)
bpy.ops.import_scene.gltf(filepath=args.input)
custom_shapes = {bone.custom_shape for obj in bpy.context.scene.objects if obj.type == 'ARMATURE'
                 for bone in obj.pose.bones if bone.custom_shape is not None}
meshes = [obj for obj in bpy.context.scene.objects
          if obj.type == 'MESH' and obj not in custom_shapes and not obj.hide_render]
if not meshes:
    raise RuntimeError('No mesh to render')
if args.material_mode == 'clay':
    clay = bpy.data.materials.new('Diagnostic clay')
    clay.use_nodes = True
    surface = clay.node_tree.nodes.get('Principled BSDF')
    surface.inputs['Base Color'].default_value = (.43, .45, .48, 1)
    surface.inputs['Metallic'].default_value = 0
    surface.inputs['Roughness'].default_value = .75
    for obj in meshes:
        obj.data.materials.clear()
        obj.data.materials.append(clay)
        for face in obj.data.polygons:
            face.material_index = 0
elif args.material_mode == 'basecolor':
    converted = {}
    for obj in meshes:
        if not obj.data.materials:
            material = bpy.data.materials.new('Default base color')
            material.use_nodes = True
            obj.data.materials.append(material)
        for slot in obj.material_slots:
            original = slot.material
            if original is None:
                raise RuntimeError('Cannot isolate base color from an empty material slot')
            if original.name not in converted:
                material = original.copy()
                material.name = 'Diagnostic base color ' + original.name
                material.use_nodes = True
                nodes, links = material.node_tree.nodes, material.node_tree.links
                source = next((node for node in nodes if node.type == 'BSDF_PRINCIPLED'), None)
                if source is None:
                    source = next((node for node in nodes if node.type == 'EMISSION'), None)
                if source is None:
                    raise RuntimeError('Cannot isolate base color from imported material ' + original.name)
                color = source.inputs['Base Color'] if source.type == 'BSDF_PRINCIPLED' else source.inputs['Color']
                emission = nodes.new('ShaderNodeEmission')
                emission.inputs['Strength'].default_value = 1
                if color.is_linked:
                    links.new(color.links[0].from_socket, emission.inputs['Color'])
                else:
                    emission.inputs['Color'].default_value = color.default_value
                output = next((node for node in nodes if node.type == 'OUTPUT_MATERIAL' and node.is_active_output), None)
                if output is None:
                    output = nodes.new('ShaderNodeOutputMaterial')
                shader = emission.outputs[0]
                alpha = source.inputs.get('Alpha')
                if alpha is not None and (alpha.is_linked or alpha.default_value < 1):
                    transparent = nodes.new('ShaderNodeBsdfTransparent')
                    mix = nodes.new('ShaderNodeMixShader')
                    if alpha.is_linked:
                        links.new(alpha.links[0].from_socket, mix.inputs[0])
                    else:
                        mix.inputs[0].default_value = alpha.default_value
                    links.new(transparent.outputs[0], mix.inputs[1])
                    links.new(shader, mix.inputs[2])
                    shader = mix.outputs[0]
                links.new(shader, output.inputs['Surface'])
                converted[original.name] = material
            slot.material = converted[original.name]
scene = bpy.context.scene
configure_cpu_render(scene)
scene.cycles.samples = 8 if scene.cycles.use_denoising else 32
scene.render.resolution_x = scene.render.resolution_y = 384
scene.render.resolution_percentage = 100
scene.render.image_settings.file_format = 'PNG'
scene.world.use_nodes = True
scene.world.node_tree.nodes['Background'].inputs['Color'].default_value = (.15, .17, .2, 1)
scene.world.node_tree.nodes['Background'].inputs['Strength'].default_value = .6
points = [obj.matrix_world @ Vector(corner) for obj in meshes for corner in obj.bound_box]
low = Vector([min(p[i] for p in points) for i in range(3)])
high = Vector([max(p[i] for p in points) for i in range(3)])
center, size = (low + high) / 2, max(high - low)
if not math.isfinite(size) or size <= 0:
    raise RuntimeError('Invalid model bounds')
for name, xyz, energy in [('key', (3, -4, 5), 1000), ('fill', (-3, -2, 2), 600), ('rim', (2, 4, 4), 900)]:
    data = bpy.data.lights.new(name, 'AREA')
    data.energy = energy * size * size / 9
    data.shape = 'DISK'
    data.size = size * 1.5
    obj = bpy.data.objects.new(name, data)
    scene.collection.objects.link(obj)
    obj.location = center + Vector(xyz) * size / 3
    obj.rotation_euler = (center - obj.location).to_track_quat('-Z', 'Y').to_euler()
camera_data = bpy.data.cameras.new('Review')
camera = bpy.data.objects.new('Review', camera_data)
scene.collection.objects.link(camera)
scene.camera = camera
camera_data.type = 'ORTHO'
camera_data.clip_end = size * 100 + 100


rendered = []


def render(name, direction, target=center, scale=size):
    camera_data.ortho_scale = scale * 1.3
    camera.location = target + Vector(direction).normalized() * size * 3
    camera.rotation_euler = (target - camera.location).to_track_quat('-Z', 'Y').to_euler()
    scene.render.filepath = str(folder / (name + '.png'))
    bpy.ops.render.render(write_still=True)
    label = name if args.material_mode == 'original' else args.material_mode + ' / ' + name
    rendered.append({'file': name + '.png', 'label': label})


for name, direction in [('front', (0, -1, .1)), ('back', (0, 1, .1)), ('left', (-1, 0, .1)),
                        ('right', (1, 0, .1)), ('top', (0, .001, 1)), ('quarter', (1, -1.4, .7))]:
    if args.material_mode == 'original' or name in {'front', 'back', 'quarter'}:
        render(name, direction)
if args.components:
    for obj in meshes:
        for other in meshes:
            other.hide_render = other != obj
        bounds = [obj.matrix_world @ Vector(corner) for corner in obj.bound_box]
        minimum = Vector([min(p[i] for p in bounds) for i in range(3)])
        maximum = Vector([max(p[i] for p in bounds) for i in range(3)])
        render(obj.name, (1, -1.4, .7), (minimum + maximum) / 2, max(maximum - minimum))

(folder / 'render-manifest.json').write_text(json.dumps({
    'material_mode': args.material_mode,
    'diagnostic_only': args.material_mode != 'original',
    'source_modified': False,
    'images': rendered,
    'denoising': bool(scene.cycles.use_denoising),
    'samples': int(scene.cycles.samples),
}, indent=2))