blender_scripts/render_objaverse.py

# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse, sys, os, math, re
import bpy
from mathutils import Vector, Matrix
import numpy as np
import cv2
import signal
from contextlib import contextmanager
from loguru import logger
from typing import Any, Callable, Dict, Generator, List, Literal, Optional, Set, Tuple
import random
class TimeoutException(Exception): pass

logger.info('Rendering started.')

@contextmanager
def time_limit(seconds):
    def signal_handler(signum, frame):
        raise TimeoutException("Timed out!")
    signal.signal(signal.SIGALRM, signal_handler)
    signal.alarm(seconds)
    try:
        yield
    finally:
        signal.alarm(0)

parser = argparse.ArgumentParser(description='Renders given obj file by rotation a camera around it.')
parser.add_argument(
    '--seed', type=int, default=0,
    help='number of views to be rendered')
parser.add_argument(
    '--views', type=int, default=4,
    help='number of views to be rendered')
parser.add_argument(
    'obj', type=str,
    help='Path to the obj file to be rendered.')
parser.add_argument(
    '--output_folder', type=str, default='/tmp',
    help='The path the output will be dumped to.')
parser.add_argument(
    '--scale', type=float, default=1,
    help='Scaling factor applied to model. Depends on size of mesh.')
parser.add_argument(
    '--format', type=str, default='PNG',
    help='Format of files generated. Either PNG or OPEN_EXR')

parser.add_argument(
    '--resolution', type=int, default=512,
    help='Resolution of the images.')
parser.add_argument(
    '--engine', type=str, default='CYCLES',
    help='Blender internal engine for rendering. E.g. CYCLES, BLENDER_EEVEE, ...')
parser.add_argument(
    '--gpu', type=int, default=0,
    help='gpu.')
parser.add_argument(
    '--animation_idx', type=int, default=0,
    help='The index of animation')

parser.add_argument(
    '--camera_option', type=str, default='fixed',
    help='Camera Options')
parser.add_argument(
    '--fixed_animation_length', type=int, default=-1,
    help='Set animation length to fixed number of framnes')
parser.add_argument(
    '--step_angle', type=int, default=3,
    help='Angle in degree for each step camera rotation')
parser.add_argument(
    '--downsample', type=int, default=1,
    help='Downsample ratio. No downsample by default')

argv = sys.argv[sys.argv.index("--") + 1:]
args = parser.parse_args(argv)


model_identifier = os.path.split(args.obj)[1].split('.')[0]
synset_idx = args.obj.split('/')[-2]

save_root = os.path.join(os.path.abspath(args.output_folder),  synset_idx, model_identifier, f'{args.animation_idx:03d}')

# Set up rendering
context = bpy.context
scene = bpy.context.scene
render = bpy.context.scene.render

render.engine = args.engine# 'BLENDER_EEVEE'
render.image_settings.color_mode = 'RGBA'  # ('RGB', 'RGBA', ...)
render.image_settings.file_format = args.format  # ('PNG', 'OPEN_EXR', 'JPEG, ...)
render.resolution_x = args.resolution
render.resolution_y = args.resolution
render.resolution_percentage = 100
bpy.context.scene.cycles.filter_width = 0.01
bpy.context.scene.render.film_transparent = True
render_depth_normal = False
bpy.context.scene.cycles.device = 'GPU'
bpy.context.scene.cycles.diffuse_bounces = 1
bpy.context.scene.cycles.glossy_bounces = 1
bpy.context.scene.cycles.transparent_max_bounces = 1
bpy.context.scene.cycles.transmission_bounces = 1
bpy.context.scene.cycles.samples = 16
bpy.context.scene.cycles.use_denoising = True
bpy.context.scene.cycles.denoiser = 'OPTIX'
bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA'
bpy.context.scene.cycles.device = 'GPU'


def enable_cuda_devices():
    prefs = bpy.context.preferences
    cprefs = prefs.addons['cycles'].preferences
    cprefs.get_devices()
    # Attempt to set GPU device types if available
    for compute_device_type in ('CUDA', 'OPENCL', 'NONE'):
        try:
            cprefs.compute_device_type = compute_device_type
            print("Compute device selected: {0}".format(compute_device_type))
            break
        except TypeError:
            pass

    # Any CUDA/OPENCL devices?
    acceleratedTypes = ['CUDA', 'OPENCL', 'OPTIX']
    acceleratedTypes = ['CUDA', 'OPENCL']
    accelerated = any(device.type in acceleratedTypes for device in cprefs.devices)
    print('Accelerated render = {0}'.format(accelerated))

    # If we have CUDA/OPENCL devices, enable only them, otherwise enable
    # all devices (assumed to be CPU)
    print(cprefs.devices)
    for idx, device in enumerate(cprefs.devices):
        device.use = (not accelerated or device.type in acceleratedTypes)# and idx == args.gpu
        print('Device enabled ({type}) = {enabled}'.format(type=device.type, enabled=device.use))
    return accelerated

enable_cuda_devices()
context.active_object.select_set(True)
bpy.ops.object.delete()

# Import textured mesh
bpy.ops.object.select_all(action='DESELECT')

try:
    with time_limit(1000):
        imported_object = bpy.ops.import_scene.gltf(filepath=args.obj, merge_vertices=True, guess_original_bind_pose=False, bone_heuristic="TEMPERANCE")
except TimeoutException as e:
    print("Timed out finished!")
    exit()


# count animated frames
animation_names = []
ending_frame_list = {}
for k in bpy.data.actions.keys():
    matched_obj_name = ''
    for obj in bpy.context.selected_objects:
        if '_'+obj.name in k and len(obj.name) > len(matched_obj_name):
            matched_obj_name = obj.name
    a_name = k.replace('_'+matched_obj_name, '')
    a = bpy.data.actions[k]
    frame_start, frame_end = map(int, a.frame_range)
    logger.info(f'{k} | frame start: {frame_start}, frame end: {frame_end} | fps: {bpy.context.scene.render.fps}')
    if a_name not in animation_names:
        animation_names.append(a_name)
        ending_frame_list[a_name] = frame_end
    else:
        ending_frame_list[a_name] = max(frame_end, ending_frame_list[a_name])
        

        
selected_a_name = animation_names[args.animation_idx]
max_frame = ending_frame_list[selected_a_name]
for obj in bpy.context.selected_objects:
    if obj.animation_data is not None:
        obj_a_name = selected_a_name+'_'+obj.name
        if obj_a_name in bpy.data.actions:
            print('Found ', obj_a_name)
            obj.animation_data.action = bpy.data.actions[obj_a_name]
        else:
            print('Miss ', obj_a_name)
            
num_frames = args.fixed_animation_length if args.fixed_animation_length != -1 else max_frame
num_frames = num_frames // args.downsample

if num_frames == 0:
    print("No animation!")
    exit()

# from https://github.com/allenai/objaverse-xl/blob/main/scripts/rendering/blender_script.py
def get_3x4_RT_matrix_from_blender(cam: bpy.types.Object):
    """Returns the 3x4 RT matrix from the given camera.

    Taken from Zero123, which in turn was taken from
    https://github.com/panmari/stanford-shapenet-renderer/blob/master/render_blender.py

    Args:
        cam (bpy.types.Object): The camera object.

    Returns:
        Matrix: The 3x4 RT matrix from the given camera.
    """
    # Use matrix_world instead to account for all constraints
    location, rotation = cam.matrix_world.decompose()[0:2]
    R_world2bcam = rotation.to_matrix().transposed()

    # Use location from matrix_world to account for constraints:
    T_world2bcam = -1 * R_world2bcam @ location

    # put into 3x4 matrix
    RT = Matrix(
        (
            R_world2bcam[0][:] + (T_world2bcam[0],),
            R_world2bcam[1][:] + (T_world2bcam[1],),
            R_world2bcam[2][:] + (T_world2bcam[2],),
        )
    )
    return RT
def _create_light(
    name: str,
    light_type: Literal["POINT", "SUN", "SPOT", "AREA"],
    location: Tuple[float, float, float],
    rotation: Tuple[float, float, float],
    energy: float,
    use_shadow: bool = False,
    specular_factor: float = 1.0,
):
    """Creates a light object.

    Args:
        name (str): Name of the light object.
        light_type (Literal["POINT", "SUN", "SPOT", "AREA"]): Type of the light.
        location (Tuple[float, float, float]): Location of the light.
        rotation (Tuple[float, float, float]): Rotation of the light.
        energy (float): Energy of the light.
        use_shadow (bool, optional): Whether to use shadows. Defaults to False.
        specular_factor (float, optional): Specular factor of the light. Defaults to 1.0.

    Returns:
        bpy.types.Object: The light object.
    """

    light_data = bpy.data.lights.new(name=name, type=light_type)
    light_object = bpy.data.objects.new(name, light_data)
    bpy.context.collection.objects.link(light_object)
    light_object.location = location
    light_object.rotation_euler = rotation
    light_data.use_shadow = use_shadow
    light_data.specular_factor = specular_factor
    light_data.energy = energy
    return light_object


def randomize_lighting() -> Dict[str, bpy.types.Object]:
    """Randomizes the lighting in the scene.

    Returns:
        Dict[str, bpy.types.Object]: Dictionary of the lights in the scene. The keys are
            "key_light", "fill_light", "rim_light", and "bottom_light".
    """

    # Clear existing lights
    bpy.ops.object.select_all(action="DESELECT")
    bpy.ops.object.select_by_type(type="LIGHT")
    bpy.ops.object.delete()

    # Create key light
    key_light = _create_light(
        name="Key_Light",
        light_type="SUN",
        location=(0, 0, 0),
        rotation=(0.785398, 0, -0.785398),
#         energy=random.choice([3, 4, 5]),
        energy=4,
    )

    # Create fill light
    fill_light = _create_light(
        name="Fill_Light",
        light_type="SUN",
        location=(0, 0, 0),
        rotation=(0.785398, 0, 2.35619),
#         energy=random.choice([2, 3, 4]),
        energy=3,
    )

    # Create rim light
    rim_light = _create_light(
        name="Rim_Light",
        light_type="SUN",
        location=(0, 0, 0),
        rotation=(-0.785398, 0, -3.92699),
#         energy=random.choice([3, 4, 5]),
        energy=4,
    )

    # Create bottom light
    bottom_light = _create_light(
        name="Bottom_Light",
        light_type="SUN",
        location=(0, 0, 0),
        rotation=(3.14159, 0, 0),
#         energy=random.choice([1, 2, 3]),
        energy=2,
    )

    return dict(
        key_light=key_light,
        fill_light=fill_light,
        rim_light=rim_light,
        bottom_light=bottom_light,
    )

def scene_bbox(
    single_obj = None, ignore_matrix = False
):
    """Returns the bounding box of the scene.

    Taken from Shap-E rendering script
    (https://github.com/openai/shap-e/blob/main/shap_e/rendering/blender/blender_script.py#L68-L82)

    Args:
        single_obj (Optional[bpy.types.Object], optional): If not None, only computes
            the bounding box for the given object. Defaults to None.
        ignore_matrix (bool, optional): Whether to ignore the object's matrix. Defaults
            to False.

    Raises:
        RuntimeError: If there are no objects in the scene.

    Returns:
        Tuple[Vector, Vector]: The minimum and maximum coordinates of the bounding box.
    """
    bbox_min = (math.inf,) * 3
    bbox_max = (-math.inf,) * 3
    found = False
    for i in range(num_frames):
        bpy.context.scene.frame_set(i * args.downsample)
        for obj in get_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 get_scene_meshes():
    """Returns all meshes in the scene.

    Yields:
        Generator[bpy.types.Object, None, None]: Generator of all meshes in the scene.
    """
    for obj in bpy.context.scene.objects.values():
        if isinstance(obj.data, (bpy.types.Mesh)):
            yield obj

def get_scene_root_objects():
    """Returns all root objects in the scene.

    Yields:
        Generator[bpy.types.Object, None, None]: Generator of all root objects in the
            scene.
    """
    for obj in bpy.context.scene.objects.values():
        if not obj.parent:
            yield obj
            
def normalize_scene():
    """Normalizes the scene by scaling and translating it to fit in a unit cube centered
    at the origin.

    Mostly taken from the Point-E / Shap-E rendering script
    (https://github.com/openai/point-e/blob/main/point_e/evals/scripts/blender_script.py#L97-L112),
    but fix for multiple root objects: (see bug report here:
    https://github.com/openai/shap-e/pull/60).

    Returns:
        None
    """
    if len(list(get_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 get_scene_root_objects():
            if obj != parent_empty:
                obj.parent = parent_empty

    bbox_min, bbox_max = scene_bbox()
    scale = 1 / max(bbox_max - bbox_min)
    logger.info(f"Scale: {scale}")
    for obj in get_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 get_scene_root_objects():
        obj.matrix_world.translation += offset
    bpy.ops.object.select_all(action="DESELECT")

    # unparent the camera
    bpy.data.objects["Camera"].parent = None

normalize_scene()

randomize_lighting()

# Place camera
cam = scene.objects['Camera']
cam.location = (0, 1.5, 0)  # radius equals to 1
cam.data.lens = 35
cam.data.sensor_width = 32

cam_constraint = cam.constraints.new(type='TRACK_TO')
cam_constraint.track_axis = 'TRACK_NEGATIVE_Z'
cam_constraint.up_axis = 'UP_Y'

cam_empty = bpy.data.objects.new("Empty", None)
cam_empty.location = (0, 0, 0)
cam.parent = cam_empty

scene.collection.objects.link(cam_empty)
context.view_layer.objects.active = cam_empty
cam_constraint.target = cam_empty

stepsize = 360.0 / args.views
rotation_mode = 'XYZ'


np.random.seed(args.seed)

if args.camera_option == "fixed":
    for scene in bpy.data.scenes:
        scene.cycles.device = 'GPU'
    
    elevation_angle = 0.
    rotation_angle = 0.

    for view_idx in range(args.views):
        img_folder = os.path.join(save_root, f'{view_idx:03d}', 'img')
        mask_folder = os.path.join(save_root, f'{view_idx:03d}', 'mask')
        camera_folder = os.path.join(save_root, f'{view_idx:03d}', 'camera')

        os.makedirs(img_folder, exist_ok=True)
        os.makedirs(mask_folder, exist_ok=True)
        os.makedirs(camera_folder, exist_ok=True)
        
        np.save(os.path.join(camera_folder, 'rotation'), np.array([rotation_angle + view_idx * stepsize for _ in range(num_frames)]))
        np.save(os.path.join(camera_folder, 'elevation'), np.array([elevation_angle for _ in range(num_frames)]))

        cam_empty.rotation_euler[2] = math.radians(rotation_angle + view_idx * stepsize)
        cam_empty.rotation_euler[0] = math.radians(elevation_angle)
        
        # save camera RT matrix
        rt_matrix = get_3x4_RT_matrix_from_blender(cam)
        rt_matrix_path = os.path.join(camera_folder, "rt_matrix.npy")
        np.save(rt_matrix_path, rt_matrix)
        for i in range(0, num_frames):
            bpy.context.scene.frame_set(i * args.downsample)
            render_file_path = os.path.join(img_folder,'%03d.png' % (i))
            scene.render.filepath = render_file_path
            bpy.ops.render.render(write_still=True)

        for i in range(0, num_frames):
            img = cv2.imread(os.path.join(img_folder, '%03d.png' % (i)), cv2.IMREAD_UNCHANGED)
            mask =  img[:, :, 3:4] / 255.0
            white_img = img[:, :, :3] * mask + np.ones_like(img[:, :, :3]) * (1 - mask) * 255
            white_img = np.clip(white_img, 0, 255)
            cv2.imwrite(os.path.join(img_folder, '%03d.jpg' % (i)), white_img)
            cv2.imwrite(os.path.join(mask_folder, '%03d.png'%(i)), img[:, :, 3])
            os.system('rm %s'%(os.path.join(img_folder, '%03d.png' % (i))))
        
elif args.camera_option == "random":
    for scene in bpy.data.scenes:
        scene.cycles.device = 'GPU'
    
    for view_idx in range(args.views):
        elevation_angle = np.random.rand(1) * 35 - 5 # [-5, 30]
        rotation_angle = np.random.rand(1) * 360
        
        img_folder = os.path.join(save_root, f'{view_idx:03d}', 'img')
        mask_folder = os.path.join(save_root, f'{view_idx:03d}', 'mask')
        camera_folder = os.path.join(save_root, f'{view_idx:03d}', 'camera')

        os.makedirs(img_folder, exist_ok=True)
        os.makedirs(mask_folder, exist_ok=True)
        os.makedirs(camera_folder, exist_ok=True)
        
        np.save(os.path.join(camera_folder, 'rotation'), np.array([rotation_angle for _ in range(num_frames)]))
        np.save(os.path.join(camera_folder, 'elevation'), np.array([elevation_angle for _ in range(num_frames)]))
        
        cam_empty.rotation_euler[2] = math.radians(rotation_angle)
        cam_empty.rotation_euler[0] = math.radians(elevation_angle)
        
        # save camera RT matrix
        rt_matrix = get_3x4_RT_matrix_from_blender(cam)
        rt_matrix_path = os.path.join(camera_folder, "rt_matrix.npy")
        np.save(rt_matrix_path, rt_matrix)
        
        for i in range(0, num_frames):
            bpy.context.scene.frame_set(i * args.downsample)
            render_file_path = os.path.join(img_folder,'%03d.png' % (i))
            scene.render.filepath = render_file_path
            bpy.ops.render.render(write_still=True)

        for i in range(0, num_frames):
            img = cv2.imread(os.path.join(img_folder, '%03d.png' % (i)), cv2.IMREAD_UNCHANGED)
            mask =  img[:, :, 3:4] / 255.0
            white_img = img[:, :, :3] * mask + np.ones_like(img[:, :, :3]) * (1 - mask) * 255
            white_img = np.clip(white_img, 0, 255)
            cv2.imwrite(os.path.join(img_folder, '%03d.jpg' % (i)), white_img)
            cv2.imwrite(os.path.join(mask_folder, '%03d.png'%(i)), img[:, :, 3])
            os.system('rm %s'%(os.path.join(img_folder, '%03d.png' % (i))))
    
else:
    raise NotImplemented