From 43a475b6176e3161196de078183354eb2a823561 Mon Sep 17 00:00:00 2001 From: huangk <43126289+huangk77@users.noreply.github.com> Date: Mon, 21 Oct 2024 15:19:31 +0800 Subject: [PATCH] Add files via upload --- render_rgbd_imagefuse.py | 290 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 290 insertions(+) create mode 100644 render_rgbd_imagefuse.py diff --git a/render_rgbd_imagefuse.py b/render_rgbd_imagefuse.py new file mode 100644 index 0000000..c6bfec2 --- /dev/null +++ b/render_rgbd_imagefuse.py @@ -0,0 +1,290 @@ +''' +MIT License +Copyright (c) 2018 Wentao Yuan +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +''' + +import bpy +import mathutils +import numpy as np +import os +import sys +import time + + +def random_pose(): + angle_x = np.random.uniform() * 2 * np.pi + angle_y = np.random.uniform() * 2 * np.pi + angle_z = np.random.uniform() * 2 * np.pi + Rx = np.array([[1, 0, 0], + [0, np.cos(angle_x), -np.sin(angle_x)], + [0, np.sin(angle_x), np.cos(angle_x)]]) + Ry = np.array([[np.cos(angle_y), 0, np.sin(angle_y)], + [0, 1, 0], + [-np.sin(angle_y), 0, np.cos(angle_y)]]) + Rz = np.array([[np.cos(angle_z), -np.sin(angle_z), 0], + [np.sin(angle_z), np.cos(angle_z), 0], + [0, 0, 1]]) + R = np.dot(Rz, np.dot(Ry, Rx)) + # Set camera pointing to the origin and 1 unit away from the origin + t = np.expand_dims(R[:, 2], 1) + pose = np.concatenate([np.concatenate([R, t], 1), [[0, 0, 0, 1]]], 0) + return pose + +def quaternion2roatation(q): + w = q[0] + x = q[1] + y = q[2] + z = q[3] + R = np.array([[1-2*y*y-2*z*z, 2*x*y-2*z*w, 2*x*z+2*y*w], + [2*x*y+2*z*w, 1-2*x*x-2*z*z, 2*y*z-2*x*w], + [2*x*z-2*y*w, 2*y*z+2*x*w, 1-2*x*x-2*y*y]]) + return R + +def pose_to_transformation_matrix(input_data): + """ + Convert pose to transformation matrix + :param position: 3D position + :param quaternion: quaternion + :return: transformation matrix + """ + quaternion = input_data[0:4] + translation = input_data[4:7] + transformation_matrix = np.zeros((4, 4)) + transformation_matrix[:3, :3] = quaternion2roatation(quaternion) + transformation_matrix[:3, 3] = translation + transformation_matrix[3, 3] = 1 + return transformation_matrix + +def setup_blender(width, height, focal_length): + # camera + camera = bpy.data.objects['Camera'] + camera.data.clip_end = 1000 + # camera.data.angle = np.arctan(width / 2 / focal_length) * 2 + camera.data.lens = focal_length + # camera.rotation_mode = 'QUATERNION' + camera.data.sensor_fit = 'AUTO' + camera.data.sensor_width = 7 + + # render layer + scene = bpy.context.scene + scene.render.filepath = 'buffer' + scene.view_layers['ViewLayer'].use_pass_z = True + scene.render.image_settings.color_depth = '8' + scene.render.image_settings.use_zbuffer = True + scene.render.resolution_percentage = 100 + scene.render.resolution_x = width + scene.render.resolution_y = height + + # this is for transparent background + scene.render.film_transparent = True + + sun_data = bpy.data.lights.new(name='Sun1', type='SUN') + sun_data.energy = 0.3 + Sun_object = bpy.data.objects.new(name='Sun1', object_data=sun_data) + Sun_object.rotation_mode = 'QUATERNION' + Sun_object.location = (0, 0, 0) + bpy.context.scene.collection.objects.link(Sun_object) + + + # # compositor nodes + # scene.use_nodes = True + # tree = scene.node_tree + # # remove nodes before create new nodes + # for n in tree.nodes: + # tree.nodes.remove(n) + # + # # create new nodes + # render_layers = tree.nodes.new('CompositorNodeRLayers') + # + # depth_file_output = tree.nodes.new('CompositorNodeOutputFile') + # depth_file_output.label = 'Depth Output' + # depth_file_output.format.file_format = 'OPEN_EXR' + # depth_file_output.format.color_depth = '32' + # depth_file_output.format.use_zbuffer = True + # depth_file_output.base_path = '' + # tree.links.new(render_layers.outputs['Depth'], depth_file_output.inputs[0]) + # + # image_file_output = tree.nodes.new('CompositorNodeOutputFile') + # image_file_output.label = 'Image' + # image_file_output.base_path = '' + # tree.links.new(render_layers.outputs['Image'], image_file_output.inputs[0]) + + # remove default cube + bpy.data.objects['Cube'].select_set(True) + bpy.ops.object.delete() + # remove default light + bpy.data.objects['Light'].select_set(True) + bpy.ops.object.delete() + + bpy.context.scene.render.engine = 'CYCLES' + # bpy.context.scene.cycles.feature_set = 'EXPERIMENTAL' + bpy.context.scene.cycles.device = 'CPU' + + # return scene, camera, Sun_object, depth_file_output, image_file_output + return scene, camera, Sun_object + + +def generate_composeitor_nodes(img_file): + scene.use_nodes = True + tree = scene.node_tree + # remove nodes before create new nodes + for n in tree.nodes: + tree.nodes.remove(n) + + # create new nodes + render_layers = tree.nodes.new('CompositorNodeRLayers') + + image_node = tree.nodes.new('CompositorNodeImage') + image_node.image = bpy.data.images.load(img_file) + + # create alpha over node + alpha_over_node = tree.nodes.new('CompositorNodeAlphaOver') + scale_node = tree.nodes.new('CompositorNodeScale') + + scale_node.space = 'RENDER_SIZE' + scale_node.frame_method = 'CROP' + + # connect nodes + tree.links.new(image_node.outputs['Image'], scale_node.inputs[0]) + tree.links.new(scale_node.outputs['Image'], alpha_over_node.inputs[1]) + tree.links.new(render_layers.outputs['Image'], alpha_over_node.inputs[2]) + + + depth_file_output = tree.nodes.new('CompositorNodeOutputFile') + depth_file_output.label = 'Depth Output' + depth_file_output.format.file_format = 'OPEN_EXR' + depth_file_output.format.color_depth = '32' + depth_file_output.format.use_zbuffer = True + depth_file_output.base_path = '' + tree.links.new(render_layers.outputs['Depth'], depth_file_output.inputs[0]) + + image_file_output = tree.nodes.new('CompositorNodeOutputFile') + image_file_output.label = 'Image' + image_file_output.base_path = '' + tree.links.new(alpha_over_node.outputs['Image'], image_file_output.inputs[0]) + + return depth_file_output, image_file_output + +if __name__ == '__main__': + # model_dir = sys.argv[-4] + # list_path = sys.argv[-3] + # output_dir = sys.argv[-2] + # num_scans = int(sys.argv[-1]) + + model_dir = './available_model/' + output_dir = './scan_data_sim_satellite_512/' + pose_dir = './pose_512.txt' + list_path = './model_list.txt' + image_file = './background/truth/' + num_scans = 10 + + nn = 1 + width = 512 * nn + height = 512 * nn + focal = 10 * nn + scene, camera, Sun = setup_blender(width, height, focal) + intrinsics = np.array([[focal, 0, width / 2], [0, focal, height / 2], [0, 0, 1]]) + + with open(os.path.join(list_path)) as file: + model_list = [line.strip() for line in file] + + pose_data = [] + with open(os.path.join(pose_dir)) as file: + for line in file: + parts = line.strip().split(',') + txt_data = [float(p) for p in parts] + pose_data.append(txt_data) + + img_file_name = os.listdir(image_file) + + pose_data = np.array(pose_data) + quaternions = pose_data[:, 0:4] + translations = pose_data[:, 4:7] + # print(pose_data[0,:]) + + open('blender.log', 'w+').close() + # os.system('rm -rf %s' % output_dir) + if not os.path.exists(output_dir): + os.makedirs(output_dir) + np.savetxt(os.path.join(output_dir, 'intrinsics.txt'), intrinsics, '%f') + + j=0 + for model_id in model_list: + start = time.time() + exr_dir = os.path.join(output_dir, model_id, 'exr') + rgb_dir = os.path.join(output_dir, model_id, 'rgb') + pose_dir = os.path.join(output_dir, model_id, 'pose') + os.makedirs(exr_dir) + os.makedirs(pose_dir) + os.makedirs(rgb_dir) + + # Redirect output to log file + old_os_out = os.dup(1) + os.close(1) + os.open('blender.log', os.O_WRONLY) + + # Import mesh model + model_path = os.path.join(model_dir, model_id+'.obj') + bpy.ops.import_scene.obj(filepath=model_path) + obj = bpy.context.selected_objects[0] + # obj.name = 'Model' + # obj.data.name = 'Model' + obj.rotation_mode = 'QUATERNION' + # + + # Rotate model by 90 degrees around x-axis (z-up => y-up) to match ShapeNet's coordinates + # bpy.ops.transform.rotate(value=-np.pi / 2, orient_axis='X') + + # Render + + for i in range(num_scans): + scene.frame_set(i) + # pose = pose_to_transformation_matrix(pose_data[i,:]) + # print(quaternions[i,:]) + Sun.rotation_quaternion = quaternions[j,:] + + img_file = os.path.join(image_file, img_file_name[i]) + depth_file_output, image_file_output = generate_composeitor_nodes(img_file) + + # camera.matrix_world = mathutils.Matrix(pose) + obj.rotation_quaternion = quaternions[j,:] + obj.location = translations[j,:] + camera.rotation_euler = (np.pi,0,0) + camera.location = (0,0,0) + j += 1 + # pose_data = pose_data[i, 0:8] + # scene.render.filepath = exr_dir + depth_file_output.file_slots[0].path = os.path.join(exr_dir, '#.exr') + image_file_output.file_slots[0].path = os.path.join(rgb_dir, '#.png') + # depth_file_output.file_slots[0].path = os.path.join(exr_dir, '#.exr') + bpy.ops.render.render(write_still=True) + # np.savetxt(os.path.join(pose_dir, '%d.txt' % i), pose, '%f') + + # Clean up + bpy.ops.object.delete() + for m in bpy.data.meshes: + bpy.data.meshes.remove(m) + for m in bpy.data.materials: + m.user_clear() + bpy.data.materials.remove(m) + + # Show time + os.close(1) + os.dup(old_os_out) + os.close(old_os_out) + print('%s done, time=%.4f sec' % (model_id, time.time() - start)) \ No newline at end of file