From 5905669dbdbb58fd4debb1295cd3d2cc182cb3b6 Mon Sep 17 00:00:00 2001 From: hofee <64160135+GitHofee@users.noreply.github.com> Date: Tue, 15 Oct 2024 11:07:17 -0500 Subject: [PATCH] format --- blender_util.py | 265 +++++++++++++++++++++++++++--------------------- 1 file changed, 148 insertions(+), 117 deletions(-) diff --git a/blender_util.py b/blender_util.py index 7e5fe11..ada8f44 100644 --- a/blender_util.py +++ b/blender_util.py @@ -1,4 +1,3 @@ - import os import json import bpy @@ -7,6 +6,7 @@ import gc import numpy as np import mathutils + class BlenderUtils: TABLE_NAME: str = "table" @@ -19,7 +19,7 @@ class BlenderUtils: @staticmethod def get_obj_path(obj_dir, name): return os.path.join(obj_dir, name, BlenderUtils.MESH_FILE_NAME) - + @staticmethod def load_obj(name, mesh_path, scale=1): print(mesh_path) @@ -34,32 +34,31 @@ class BlenderUtils: @staticmethod def get_obj(name): return bpy.data.objects.get(name) - - + @staticmethod def get_obj_pose(name): obj = BlenderUtils.get_obj(name) return np.asarray(obj.matrix_world) - - - + @staticmethod def add_plane(name, location, orientation, size=10): - bpy.ops.mesh.primitive_plane_add(size=size,location=location) + bpy.ops.mesh.primitive_plane_add(size=size, location=location) plane = bpy.context.selected_objects[-1] plane.name = name plane.rotation_euler = orientation bpy.ops.rigidbody.object_add() - bpy.context.object.rigid_body.type = 'PASSIVE' + bpy.context.object.rigid_body.type = "PASSIVE" @staticmethod def add_table(table_model_path): - table = BlenderUtils.load_obj(BlenderUtils.TABLE_NAME, table_model_path, scale=0.01) + table = BlenderUtils.load_obj( + BlenderUtils.TABLE_NAME, table_model_path, scale=0.01 + ) bpy.ops.rigidbody.object_add() - bpy.context.object.rigid_body.type = 'PASSIVE' - + bpy.context.object.rigid_body.type = "PASSIVE" + mat = bpy.data.materials.new(name="TableYellowMaterial") - mat.diffuse_color = (1.0, 1.0, 0.0, 1.0) + mat.diffuse_color = (1.0, 1.0, 0.0, 1.0) if len(table.data.materials) > 0: table.data.materials[0] = mat else: @@ -67,107 +66,124 @@ class BlenderUtils: @staticmethod def setup_scene(init_light_and_camera_config, table_model_path, binocular_vision): - bpy.context.scene.render.engine = 'BLENDER_EEVEE_NEXT' + bpy.context.scene.render.engine = "BLENDER_EEVEE_NEXT" bpy.context.scene.display.shading.show_xray = False bpy.context.scene.display.shading.use_dof = False - bpy.context.scene.display.render_aa = 'OFF' - bpy.context.scene.view_settings.view_transform = 'Standard' + bpy.context.scene.display.render_aa = "OFF" + bpy.context.scene.view_settings.view_transform = "Standard" - bpy.context.scene.eevee.use_ssr = False # 关闭屏幕空间反射 - bpy.context.scene.eevee.use_bloom = False # 关闭辉光 - bpy.context.scene.eevee.use_gtao = False # 关闭环境光遮蔽 - bpy.context.scene.eevee.use_soft_shadows = False # 关闭软阴影 + bpy.context.scene.eevee.use_ssr = False # 关闭屏幕空间反射 + bpy.context.scene.eevee.use_bloom = False # 关闭辉光 + bpy.context.scene.eevee.use_gtao = False # 关闭环境光遮蔽 + bpy.context.scene.eevee.use_soft_shadows = False # 关闭软阴影 bpy.context.scene.eevee.use_shadows = False # 关闭所有阴影 bpy.context.scene.world.use_nodes = False # 如果你不需要环境光，关闭环境节点 - - #bpy.context.scene.eevee.use_sss = False # 关闭次表面散射 + # bpy.context.scene.eevee.use_sss = False # 关闭次表面散射 # 2. 设置最低的采样数 bpy.context.scene.eevee.taa_render_samples = 1 bpy.context.scene.eevee.taa_samples = 1 - BlenderUtils.init_light_and_camera(init_light_and_camera_config, binocular_vision) - - BlenderUtils.add_plane("plane_floor", location=(0,0,0), orientation=(0,0,0)) - BlenderUtils.add_plane("plane_ceil", location=(0,0,10), orientation=(0,0,0)) - BlenderUtils.add_plane("plane_wall_1", location=(5,0,5), orientation=(0,np.pi/2,0)) - BlenderUtils.add_plane("plane_wall_2", location=(-5,0,5), orientation=(0,np.pi/2,0)) - BlenderUtils.add_plane("plane_wall_3", location=(0,5,5), orientation=(np.pi/2,0,0)) - BlenderUtils.add_plane("plane_wall_4", location=(0,-5,5), orientation=(np.pi/2,0,0)) + BlenderUtils.init_light_and_camera( + init_light_and_camera_config, binocular_vision + ) + + BlenderUtils.add_plane("plane_floor", location=(0, 0, 0), orientation=(0, 0, 0)) + BlenderUtils.add_plane("plane_ceil", location=(0, 0, 10), orientation=(0, 0, 0)) + BlenderUtils.add_plane( + "plane_wall_1", location=(5, 0, 5), orientation=(0, np.pi / 2, 0) + ) + BlenderUtils.add_plane( + "plane_wall_2", location=(-5, 0, 5), orientation=(0, np.pi / 2, 0) + ) + BlenderUtils.add_plane( + "plane_wall_3", location=(0, 5, 5), orientation=(np.pi / 2, 0, 0) + ) + BlenderUtils.add_plane( + "plane_wall_4", location=(0, -5, 5), orientation=(np.pi / 2, 0, 0) + ) BlenderUtils.add_table(table_model_path) - + @staticmethod def set_light_params(light, config): light.location = config["location"] light.rotation_euler = config["orientation"] - if light.type == 'SUN': + if light.type == "SUN": light.data.energy = config["power"] - elif light.type == 'POINT': + elif light.type == "POINT": light.data.energy = config["power"] - + @staticmethod def set_camera_params(camera, config, binocular_vision): - + camera_object = bpy.data.objects.new(BlenderUtils.CAMERA_OBJECT_NAME, None) bpy.context.collection.objects.link(camera_object) cameras = [bpy.data.objects.get("Camera")] - camera.location = [0,0,0] - camera.rotation_euler = [0,0,0] + camera.location = [0, 0, 0] + camera.rotation_euler = [0, 0, 0] camera.parent = camera_object if binocular_vision: left_camera = cameras[0] right_camera = left_camera.copy() right_camera.name = BlenderUtils.CAMERA_RIGHT_NAME - right_camera.data = left_camera.data.copy() + right_camera.data = left_camera.data.copy() right_camera.data.name = BlenderUtils.CAMERA_RIGHT_NAME bpy.context.collection.objects.link(right_camera) right_camera.parent = camera_object - right_camera.location = [config["eye_distance"]/2, 0, 0] - left_camera.location = [-config["eye_distance"]/2, 0, 0] - binocular_angle = config["eye_angle"] - half_angle = np.radians(binocular_angle / 2) + right_camera.location = [config["eye_distance"] / 2, 0, 0] + left_camera.location = [-config["eye_distance"] / 2, 0, 0] + binocular_angle = config["eye_angle"] + half_angle = np.radians(binocular_angle / 2) - left_camera.rotation_euler[1] = -half_angle + left_camera.rotation_euler[1] = -half_angle right_camera.rotation_euler[1] = half_angle cameras.append(right_camera) - + for camera in cameras: camera.data.clip_start = config["near_plane"] camera.data.clip_end = config["far_plane"] - + bpy.context.scene.render.resolution_x = config["resolution"][0] bpy.context.scene.render.resolution_y = config["resolution"][1] sensor_height = 24.0 - focal_length = sensor_height / (2 * np.tan(np.radians(config["fov_vertical"]) / 2)) + focal_length = sensor_height / ( + 2 * np.tan(np.radians(config["fov_vertical"]) / 2) + ) camera.data.lens = focal_length - camera.data.sensor_width = sensor_height * config["resolution"][0] / config["resolution"][1] + camera.data.sensor_width = ( + sensor_height * config["resolution"][0] / config["resolution"][1] + ) camera.data.sensor_height = sensor_height - - @staticmethod + + @staticmethod def init_light_and_camera(init_light_and_camera_config, binocular_vision): - + camera = BlenderUtils.get_obj(BlenderUtils.CAMERA_NAME) - BlenderUtils.set_camera_params(camera, init_light_and_camera_config[BlenderUtils.CAMERA_NAME], binocular_vision) - + BlenderUtils.set_camera_params( + camera, + init_light_and_camera_config[BlenderUtils.CAMERA_NAME], + binocular_vision, + ) + @staticmethod def get_obj_diag(name): obj = BlenderUtils.get_obj(name) return np.linalg.norm(obj.dimensions) - - @staticmethod + + @staticmethod def matrix_to_blender_pose(matrix): location = matrix[:3, 3] rotation_matrix = matrix[:3, :3] rotation_matrix_blender = mathutils.Matrix(rotation_matrix.tolist()) rotation_euler = rotation_matrix_blender.to_euler() return location, rotation_euler - + @staticmethod def set_camera_at(pose): camera = BlenderUtils.get_obj(BlenderUtils.CAMERA_OBJECT_NAME) location, rotation_euler = BlenderUtils.matrix_to_blender_pose(pose) - + camera.location = location camera.rotation_euler = rotation_euler @@ -177,9 +193,11 @@ class BlenderUtils: vertices_world = [obj.matrix_world @ v for v in vertices] min_z = min([v.z for v in vertices_world]) return min_z - + @staticmethod - def render_and_save(output_dir, file_name, binocular_vision=False, target_object=None): + def render_and_save( + output_dir, file_name, binocular_vision=False, target_object=None + ): target_cameras = [BlenderUtils.CAMERA_NAME] if binocular_vision: target_cameras.append(BlenderUtils.CAMERA_RIGHT_NAME) @@ -190,14 +208,15 @@ class BlenderUtils: cam_suffix = "L" if cam_name == BlenderUtils.CAMERA_NAME else "R" scene = bpy.context.scene scene.render.filepath = "" - - + mask_dir = os.path.join(output_dir, "mask") if not os.path.exists(mask_dir): os.makedirs(mask_dir) - - scene.render.filepath = os.path.join(output_dir, mask_dir, f"{file_name}_{cam_suffix}.png") - scene.render.image_settings.color_depth = '8' + + scene.render.filepath = os.path.join( + output_dir, mask_dir, f"{file_name}_{cam_suffix}.png" + ) + scene.render.image_settings.color_depth = "8" scene.render.resolution_percentage = 100 scene.render.use_overwrite = False scene.render.use_file_extension = False @@ -209,91 +228,89 @@ class BlenderUtils: for node in tree.nodes: tree.nodes.remove(node) - rl = tree.nodes.new('CompositorNodeRLayers') + rl = tree.nodes.new("CompositorNodeRLayers") - map_range = tree.nodes.new('CompositorNodeMapRange') - map_range.inputs['From Min'].default_value = 0.01 - map_range.inputs['From Max'].default_value = 5 - map_range.inputs['To Min'].default_value = 0 - map_range.inputs['To Max'].default_value = 1 - tree.links.new(rl.outputs['Depth'], map_range.inputs[0]) + map_range = tree.nodes.new("CompositorNodeMapRange") + map_range.inputs["From Min"].default_value = 0.01 + map_range.inputs["From Max"].default_value = 5 + map_range.inputs["To Min"].default_value = 0 + map_range.inputs["To Max"].default_value = 1 + tree.links.new(rl.outputs["Depth"], map_range.inputs[0]) - output_depth = tree.nodes.new('CompositorNodeOutputFile') + output_depth = tree.nodes.new("CompositorNodeOutputFile") depth_dir = os.path.join(output_dir, "depth") if not os.path.exists(depth_dir): os.makedirs(depth_dir) output_depth.base_path = depth_dir output_depth.file_slots[0].path = f"{file_name}_{cam_suffix}.####" - output_depth.format.file_format = 'PNG' - output_depth.format.color_mode = 'BW' - output_depth.format.color_depth = '16' + output_depth.format.file_format = "PNG" + output_depth.format.color_mode = "BW" + output_depth.format.color_depth = "16" - # 创建 Separate XYZ 节点来分离法线的 X, Y, Z 分量 - separate_xyz = tree.nodes.new('CompositorNodeSeparateXYZ') + # 创建 Separate XYZ 节点来分离法线的 X, Y, Z 分量 + separate_xyz = tree.nodes.new("CompositorNodeSeparateXYZ") # 将法线向量连接到 Separate XYZ 节点 - tree.links.new(rl.outputs['Normal'], separate_xyz.inputs[0]) + tree.links.new(rl.outputs["Normal"], separate_xyz.inputs[0]) # 创建 Map Range 节点来分别映射 X, Y, Z 分量 - map_range_x = tree.nodes.new('CompositorNodeMapRange') - map_range_y = tree.nodes.new('CompositorNodeMapRange') - map_range_z = tree.nodes.new('CompositorNodeMapRange') + map_range_x = tree.nodes.new("CompositorNodeMapRange") + map_range_y = tree.nodes.new("CompositorNodeMapRange") + map_range_z = tree.nodes.new("CompositorNodeMapRange") # 设置映射范围 for map_range in [map_range_x, map_range_y, map_range_z]: - map_range.inputs['From Min'].default_value = -1 - map_range.inputs['From Max'].default_value = 1 - map_range.inputs['To Min'].default_value = 0 - map_range.inputs['To Max'].default_value = 1 + map_range.inputs["From Min"].default_value = -1 + map_range.inputs["From Max"].default_value = 1 + map_range.inputs["To Min"].default_value = 0 + map_range.inputs["To Max"].default_value = 1 # 分别连接到法线的 X, Y, Z 输出 - tree.links.new(separate_xyz.outputs['X'], map_range_x.inputs[0]) - tree.links.new(separate_xyz.outputs['Y'], map_range_y.inputs[0]) - tree.links.new(separate_xyz.outputs['Z'], map_range_z.inputs[0]) + tree.links.new(separate_xyz.outputs["X"], map_range_x.inputs[0]) + tree.links.new(separate_xyz.outputs["Y"], map_range_y.inputs[0]) + tree.links.new(separate_xyz.outputs["Z"], map_range_z.inputs[0]) # 合并 X, Y, Z 分量到一个 RGB 输出 - combine_rgb = tree.nodes.new('CompositorNodeCombineXYZ') - tree.links.new(map_range_x.outputs[0], combine_rgb.inputs['X']) - tree.links.new(map_range_y.outputs[0], combine_rgb.inputs['Y']) - tree.links.new(map_range_z.outputs[0], combine_rgb.inputs['Z']) + combine_rgb = tree.nodes.new("CompositorNodeCombineXYZ") + tree.links.new(map_range_x.outputs[0], combine_rgb.inputs["X"]) + tree.links.new(map_range_y.outputs[0], combine_rgb.inputs["Y"]) + tree.links.new(map_range_z.outputs[0], combine_rgb.inputs["Z"]) # 输出到文件 - output_normal = tree.nodes.new('CompositorNodeOutputFile') + output_normal = tree.nodes.new("CompositorNodeOutputFile") normal_dir = os.path.join(output_dir, "normal") if not os.path.exists(normal_dir): os.makedirs(normal_dir) output_normal.base_path = normal_dir output_normal.file_slots[0].path = f"{file_name}_{cam_suffix}.####" - output_normal.format.file_format = 'PNG' - output_normal.format.color_mode = 'RGB' - output_normal.format.color_depth = '8' + output_normal.format.file_format = "PNG" + output_normal.format.color_mode = "RGB" + output_normal.format.color_depth = "8" tree.links.new(combine_rgb.outputs[0], output_normal.inputs[0]) bpy.ops.render.render(write_still=True) - - - + msg = "success" return msg - + @staticmethod def save_cam_params(scene_dir, idx, binocular_vision=False): camera = BlenderUtils.get_obj(BlenderUtils.CAMERA_NAME) extrinsic = np.array(camera.matrix_world) cam_data = camera.data - focal_length = cam_data.lens - sensor_width = cam_data.sensor_width - sensor_height = cam_data.sensor_height - resolution_x = bpy.context.scene.render.resolution_x - resolution_y = bpy.context.scene.render.resolution_y + focal_length = cam_data.lens + sensor_width = cam_data.sensor_width + sensor_height = cam_data.sensor_height + resolution_x = bpy.context.scene.render.resolution_x + resolution_y = bpy.context.scene.render.resolution_y intrinsic = np.zeros((3, 3)) intrinsic[0, 0] = focal_length * resolution_x / sensor_width # fx - intrinsic[1, 1] = focal_length * resolution_y / sensor_height # fy + intrinsic[1, 1] = focal_length * resolution_y / sensor_height # fy intrinsic[0, 2] = resolution_x / 2.0 # cx intrinsic[1, 2] = resolution_y / 2.0 # cy intrinsic[2, 2] = 1.0 - cam_object = BlenderUtils.get_obj(BlenderUtils.CAMERA_OBJECT_NAME) + cam_object = BlenderUtils.get_obj(BlenderUtils.CAMERA_OBJECT_NAME) extrinsic_cam_object = np.array(cam_object.matrix_world) data = { "extrinsic": extrinsic.tolist(), @@ -305,21 +322,28 @@ class BlenderUtils: if binocular_vision: right_camera = BlenderUtils.get_obj(BlenderUtils.CAMERA_RIGHT_NAME) extrinsic_right = np.array(right_camera.matrix_world) - print("result:",extrinsic_right) - + print("result:", extrinsic_right) + data["extrinsic_R"] = extrinsic_right.tolist() - + cam_params_dir = os.path.join(scene_dir, "camera_params") if not os.path.exists(cam_params_dir): os.makedirs(cam_params_dir) - cam_params_path = os.path.join(cam_params_dir, f"{idx}.json") + cam_params_path = os.path.join(cam_params_dir, f"{idx}.json") with open(cam_params_path, "w") as f: json.dump(data, f, indent=4) @staticmethod def reset_objects_and_platform(): all_objects = bpy.data.objects - keep_objects = {"plane_floor", "plane_ceil", "plane_wall_1", "plane_wall_2", "plane_wall_3", "plane_wall_4"} + keep_objects = { + "plane_floor", + "plane_ceil", + "plane_wall_1", + "plane_wall_2", + "plane_wall_3", + "plane_wall_4", + } keep_objects.add(BlenderUtils.CAMERA_OBJECT_NAME) keep_objects.add(BlenderUtils.CAMERA_NAME) keep_objects.add(BlenderUtils.CAMERA_RIGHT_NAME) @@ -345,18 +369,28 @@ class BlenderUtils: @staticmethod def save_scene_info(scene_root_dir, display_table_config, target_name): all_objects = bpy.data.objects - no_save_objects = {"plane_floor", "plane_ceil", "plane_wall_1", "plane_wall_2", "plane_wall_3", "plane_wall_4"} + no_save_objects = { + "plane_floor", + "plane_ceil", + "plane_wall_1", + "plane_wall_2", + "plane_wall_3", + "plane_wall_4", + } no_save_objects.add(BlenderUtils.CAMERA_OBJECT_NAME) no_save_objects.add(BlenderUtils.CAMERA_NAME) no_save_objects.add(BlenderUtils.CAMERA_RIGHT_NAME) no_save_objects.add(BlenderUtils.TABLE_NAME) scene_info = {} for obj in all_objects: - if obj.name not in no_save_objects and obj.name != BlenderUtils.DISPLAY_TABLE_NAME: + if ( + obj.name not in no_save_objects + and obj.name != BlenderUtils.DISPLAY_TABLE_NAME + ): obj_info = { "location": list(obj.location), "rotation_euler": list(obj.rotation_euler), - "scale": list(obj.scale) + "scale": list(obj.scale), } scene_info[obj.name] = obj_info scene_info[BlenderUtils.DISPLAY_TABLE_NAME] = display_table_config @@ -364,6 +398,3 @@ class BlenderUtils: scene_info_path = os.path.join(scene_root_dir, "scene_info.json") with open(scene_info_path, "w") as outfile: json.dump(scene_info, outfile) - - - \ No newline at end of file