import os import shutil import bpy import mathutils import numpy as np import OpenEXR import Imath from scipy.spatial.transform import Rotation as scipy_Rotation from evalmde.utils.constants import VALID_DEPTH_LB, VALID_DEPTH_UB from evalmde.utils.common import pathlib_file from evalmde.utils.depth import get_depth_valid from evalmde.utils.image import imread_rgb, imwrite_rgb def bpy_set_tmp_dir(tmp_dir): tmp_dir = pathlib_file(tmp_dir) tmp_dir.mkdir(parents=True, exist_ok=True) bpy.context.preferences.filepaths.temporary_directory = str(tmp_dir) def bpy_create_cam(cam_name, cam_pose, fx, fy, cx, cy, w, h): cam_data = bpy.data.cameras.new(name=cam_name) cam_data.sensor_height = cam_data.sensor_width * h / w cam_data.lens = (fx + fy) / 2 * cam_data.sensor_width / w cam_data.shift_x = (w / 2 - cx) / w cam_data.shift_y = (cy - h / 2) / h cam_object = bpy.data.objects.new(cam_name, cam_data) bpy.context.collection.objects.link(cam_object) cam_object.matrix_world = mathutils.Matrix([cam_pose[0], cam_pose[1], cam_pose[2], cam_pose[3]]) return cam_object def bpy_add_ambient_light(energy=1.0): world = bpy.data.worlds.new("AmbientWorld") bpy.context.scene.world = world world.use_nodes = True bg = world.node_tree.nodes["Background"] bg.inputs[0].default_value = (1, 1, 1, 1) bg.inputs[1].default_value = energy def bpy_enable_gpu(device_type="CUDA"): prefs = bpy.context.preferences cprefs = prefs.addons['cycles'].preferences cprefs.compute_device_type = device_type # "CUDA", "OPTIX", "METAL", "HIP" cprefs.get_devices() # Initialize devices for device in cprefs.devices: device.use = True bpy.context.scene.cycles.device = 'GPU' def bpy_setup_rgb_render(): bpy.context.scene.use_nodes = True tree = bpy.context.scene.node_tree tree.nodes.clear() render_layers = tree.nodes.new(type='CompositorNodeRLayers') rgb_output = tree.nodes.new(type='CompositorNodeOutputFile') rgb_output.label = 'RGB Output' rgb_output.base_path = '' rgb_output.format.file_format = 'PNG' rgb_output.file_slots[0].use_node_format = True rgb_output.file_slots[0].save_as_render = True tree.links.new(render_layers.outputs['Image'], rgb_output.inputs[0]) return rgb_output def bpy_setup_depth_render(): bpy.context.scene.use_nodes = True tree = bpy.context.scene.node_tree tree.nodes.clear() render_layers = tree.nodes.new(type='CompositorNodeRLayers') depth_output = tree.nodes.new(type='CompositorNodeOutputFile') depth_output.label = 'Depth Output' depth_output.base_path = '' depth_output.format.file_format = 'OPEN_EXR' depth_output.file_slots[0].use_node_format = True depth_output.file_slots[0].save_as_render = True bpy.context.view_layer.use_pass_z = True bpy.context.scene.view_layers["ViewLayer"].use_pass_z = True tree.links.new(render_layers.outputs['Depth'], depth_output.inputs[0]) return depth_output def bpy_setup_rgbd_render(): bpy.context.scene.use_nodes = True tree = bpy.context.scene.node_tree tree.nodes.clear() # Add Render Layers node to get passes render_layers = tree.nodes.new(type='CompositorNodeRLayers') # Add Output File node to save the EXR depth_output = tree.nodes.new(type='CompositorNodeOutputFile') depth_output.label = 'Depth Output' depth_output.base_path = '' depth_output.format.file_format = 'OPEN_EXR' depth_output.file_slots[0].use_node_format = True depth_output.file_slots[0].save_as_render = True # Output for RGB rgb_output = tree.nodes.new(type='CompositorNodeOutputFile') rgb_output.label = 'RGB Output' rgb_output.base_path = '' rgb_output.format.file_format = 'PNG' rgb_output.file_slots[0].use_node_format = True rgb_output.file_slots[0].save_as_render = True # Enable the Z (Depth) pass bpy.context.view_layer.use_pass_z = True bpy.context.scene.view_layers["ViewLayer"].use_pass_z = True # Link the depth pass output from the render layers node tree.links.new(render_layers.outputs['Depth'], depth_output.inputs[0]) tree.links.new(render_layers.outputs['Image'], rgb_output.inputs[0]) return depth_output, rgb_output def save_depth_from_exr(filepath, h, w): exr_file = OpenEXR.InputFile(filepath) dw = exr_file.header()['dataWindow'] size = (dw.max.x - dw.min.x + 1, dw.max.y - dw.min.y + 1) assert size == (w, h), f"Expected {(w, h)}, got {size}" pt = Imath.PixelType(Imath.PixelType.FLOAT) channels = exr_file.channels(["R", "G", "B"], pt) depth = [np.frombuffer(c, dtype=np.float32).reshape(size[1], size[0]) for c in channels] assert np.all(depth[0] == depth[1]) and np.all(depth[0] == depth[2]) return depth[0] def bpy_create_directional_light(src: np.ndarray, dst: np.ndarray, energy=5.0, name='Sun'): light_data = bpy.data.lights.new(name=name, type='SUN') light_data.energy = energy light_obj = bpy.data.objects.new(name=name, object_data=light_data) bpy.context.collection.objects.link(light_obj) light_obj.location = (float(src[0]), float(src[1]), float(src[2])) direction = dst - src rot_axis = np.cross(np.array([0, 0, -1.]), direction) if np.linalg.norm(rot_axis) < 1e-5: rot_axis = np.array([1., 0, 0]) rot_axis /= np.linalg.norm(rot_axis) rot_ang = np.arccos(np.clip(((direction / np.linalg.norm(direction)) * np.array([0, 0, -1.])).sum(), -1, 1)) rot_euler = scipy_Rotation.from_rotvec(rot_ang * rot_axis, degrees=False).as_euler('xyz', degrees=False) light_obj.rotation_euler = (float(rot_euler[0]), float(rot_euler[1]), float(rot_euler[2])) def bpy_render_rgb(cam_object, h, w, num_sample, rgb_node, output_root, out_name): bpy.context.scene.cycles.samples = num_sample bpy.context.scene.render.resolution_x = w bpy.context.scene.render.resolution_y = h bpy.context.scene.camera = cam_object rgb_node.base_path = str(output_root) rgb_node.file_slots[0].path = f"image_{out_name}-" bpy.ops.render.render(write_still=True) def bpy_render_rgbd(cam_object, h, w, num_sample, depth_node, rgb_node, output_root, out_name): bpy.context.scene.cycles.samples = num_sample bpy.context.scene.render.resolution_x = w bpy.context.scene.render.resolution_y = h bpy.context.scene.camera = cam_object depth_node.base_path = str(output_root) depth_node.file_slots[0].path = f"depth_{out_name}_" rgb_node.base_path = str(output_root) rgb_node.file_slots[0].path = f"image_{out_name}-" bpy.ops.render.render(write_still=True) exr_path = os.path.join(str(output_root), f"depth_{out_name}_0001.exr") depth_np = save_depth_from_exr(exr_path, h, w) np.save(os.path.join(str(output_root), f"depth_{out_name}.npy"), depth_np) os.remove(exr_path) def bpy_render_rgb_and_filter_invalid(cam_object, h, w, num_sample, depth_node, rgb_node, output_root, out_name, bkg_color, valid_depth_lb=VALID_DEPTH_LB, valid_depth_ub=VALID_DEPTH_UB, save_depth=False): ''' :param cam_object: :param h: :param w: :param num_sample: :param depth_node: :param rgb_node: :param output_root: :param out_name: :param bkg_color: list of 3 integers, [0, 255] :param valid_depth_lb: :param valid_depth_ub: :param save_depth: :return: ''' bpy_render_rgbd(cam_object, h, w, num_sample, depth_node, rgb_node, output_root, out_name) img_f = pathlib_file(os.path.join(str(output_root), f"image_{out_name}-0001.png")) depth_f = pathlib_file(os.path.join(str(output_root), f"depth_{out_name}.npy")) output_root = pathlib_file(output_root) tmp_dir = output_root.parent / f'{output_root.name}__tmp' tmp_dir.mkdir(parents=True, exist_ok=True) img = imread_rgb(img_f) depth = np.load(depth_f) shutil.move(img_f, tmp_dir / img_f.name) if not save_depth: shutil.move(depth_f, tmp_dir / depth_f.name) img[~get_depth_valid(depth, valid_depth_lb, valid_depth_ub)] = np.array(bkg_color) imwrite_rgb(output_root / f'image_{out_name}.png', img) os.remove(tmp_dir / img_f.name) os.remove(tmp_dir / depth_f.name)