import os import json import bpy import time import gc import numpy as np import mathutils class CADBlenderUtils: TABLE_NAME: str = "table" CAMERA_NAME: str = "Camera" CAMERA_RIGHT_NAME: str = "CameraRight" CAMERA_OBJECT_NAME: str = "CameraObject" DISPLAY_TABLE_NAME: str = "display_table" MESH_FILE_NAME: str = "mesh.obj" @staticmethod def get_obj_path(obj_dir, name): return os.path.join(obj_dir, name, CADBlenderUtils.MESH_FILE_NAME) @staticmethod def load_obj(name, mesh_path, scale=1): print(mesh_path) bpy.ops.wm.obj_import(filepath=mesh_path) loaded_object = bpy.context.selected_objects[-1] loaded_object.name = name loaded_object.data.name = name loaded_object.scale = (scale, scale, scale) bpy.ops.rigidbody.object_add() return loaded_object @staticmethod def get_obj(name): return bpy.data.objects.get(name) @staticmethod def get_obj_pose(name): obj = CADBlenderUtils.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) 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" @staticmethod def setup_scene(init_light_and_camera_config, table_model_path, binocular_vision): bpy.context.scene.render.engine = "BLENDER_EEVEE" 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.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 # 关闭次表面散射 # 2. 设置最低的采样数 bpy.context.scene.eevee.taa_render_samples = 1 bpy.context.scene.eevee.taa_samples = 1 CADBlenderUtils.init_light_and_camera( init_light_and_camera_config, binocular_vision ) @staticmethod def set_camera_params(camera, config, binocular_vision): camera_object = bpy.data.objects.new(CADBlenderUtils.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.parent = camera_object if binocular_vision: left_camera = cameras[0] right_camera = left_camera.copy() right_camera.name = CADBlenderUtils.CAMERA_RIGHT_NAME right_camera.data = left_camera.data.copy() right_camera.data.name = CADBlenderUtils.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) 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) ) camera.data.lens = focal_length camera.data.sensor_width = ( sensor_height * config["resolution"][0] / config["resolution"][1] ) camera.data.sensor_height = sensor_height @staticmethod def init_light_and_camera(init_light_and_camera_config, binocular_vision): camera = CADBlenderUtils.get_obj(CADBlenderUtils.CAMERA_NAME) CADBlenderUtils.set_camera_params( camera, init_light_and_camera_config[CADBlenderUtils.CAMERA_NAME], binocular_vision, ) @staticmethod def get_obj_diag(name): obj = CADBlenderUtils.get_obj(name) return np.linalg.norm(obj.dimensions) @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 = CADBlenderUtils.get_obj(CADBlenderUtils.CAMERA_OBJECT_NAME) location, rotation_euler = CADBlenderUtils.matrix_to_blender_pose(pose) camera.location = location camera.rotation_euler = rotation_euler @staticmethod def get_object_bottom_z(obj): vertices = [v.co for v in obj.data.vertices] 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_normal_and_depth( output_dir, file_name, binocular_vision=False, target_object=None ): # use pass z bpy.context.scene.view_layers["ViewLayer"].use_pass_z = True target_cameras = [CADBlenderUtils.CAMERA_NAME] if binocular_vision: target_cameras.append(CADBlenderUtils.CAMERA_RIGHT_NAME) for cam_name in target_cameras: bpy.context.scene.camera = CADBlenderUtils.get_obj(cam_name) cam_suffix = "L" if cam_name == CADBlenderUtils.CAMERA_NAME else "R" scene = bpy.context.scene scene.render.filepath = "" mask_dir = os.path.join(output_dir, "normal") 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}.exr" ) scene.render.image_settings.file_format = "OPEN_EXR" scene.render.image_settings.color_mode = "RGB" bpy.context.scene.view_settings.view_transform = "Raw" scene.render.image_settings.color_depth = "16" bpy.context.scene.render.filter_size = 1.5 scene.render.resolution_percentage = 100 scene.render.use_overwrite = False scene.render.use_file_extension = False scene.render.use_placeholder = False scene.use_nodes = True tree = scene.node_tree for node in tree.nodes: tree.nodes.remove(node) 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]) 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" tree.links.new(map_range.outputs[0], output_depth.inputs[0]) bpy.ops.render.render(write_still=True) msg = "success" return msg @staticmethod def render_mask( output_dir, file_name, binocular_vision=False, target_object=None ): target_cameras = [CADBlenderUtils.CAMERA_NAME] if binocular_vision: target_cameras.append(CADBlenderUtils.CAMERA_RIGHT_NAME) for cam_name in target_cameras: bpy.context.scene.camera = CADBlenderUtils.get_obj(cam_name) cam_suffix = "L" if cam_name == CADBlenderUtils.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.resolution_percentage = 100 scene.render.use_overwrite = False scene.render.use_file_extension = False scene.render.use_placeholder = False bpy.ops.render.render(write_still=True) msg = "success" return msg @staticmethod def save_cam_params(scene_dir, idx, binocular_vision=False): camera = CADBlenderUtils.get_obj(CADBlenderUtils.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 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[0, 2] = resolution_x / 2.0 # cx intrinsic[1, 2] = resolution_y / 2.0 # cy intrinsic[2, 2] = 1.0 cam_object = CADBlenderUtils.get_obj(CADBlenderUtils.CAMERA_OBJECT_NAME) extrinsic_cam_object = np.array(cam_object.matrix_world) data = { "extrinsic": extrinsic.tolist(), "extrinsic_cam_object": extrinsic_cam_object.tolist(), "intrinsic": intrinsic.tolist(), "far_plane": camera.data.clip_end, "near_plane": camera.data.clip_start, } if binocular_vision: right_camera = CADBlenderUtils.get_obj(CADBlenderUtils.CAMERA_RIGHT_NAME) extrinsic_right = np.array(right_camera.matrix_world) 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") 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.add(CADBlenderUtils.CAMERA_OBJECT_NAME) keep_objects.add(CADBlenderUtils.CAMERA_NAME) keep_objects.add(CADBlenderUtils.CAMERA_RIGHT_NAME) keep_objects.add(CADBlenderUtils.TABLE_NAME) for obj in all_objects: if obj.name not in keep_objects: bpy.data.objects.remove(obj, do_unlink=True) for block in bpy.data.meshes: if block.users == 0: bpy.data.meshes.remove(block) for block in bpy.data.materials: if block.users == 0: bpy.data.materials.remove(block) for block in bpy.data.images: if block.users == 0: bpy.data.images.remove(block) gc.collect() bpy.context.scene.frame_set(0) @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.add(CADBlenderUtils.CAMERA_OBJECT_NAME) no_save_objects.add(CADBlenderUtils.CAMERA_NAME) no_save_objects.add(CADBlenderUtils.CAMERA_RIGHT_NAME) no_save_objects.add(CADBlenderUtils.TABLE_NAME) scene_info = {} for obj in all_objects: if ( obj.name not in no_save_objects and obj.name != CADBlenderUtils.DISPLAY_TABLE_NAME ): obj_info = { "location": list(obj.location), "rotation_euler": list(obj.rotation_euler), "scale": list(obj.scale), } scene_info[obj.name] = obj_info scene_info[CADBlenderUtils.DISPLAY_TABLE_NAME] = display_table_config scene_info["target_name"] = target_name 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) @staticmethod def save_blend(scene_root_dir): blend_path = os.path.join(scene_root_dir, "scene.blend") bpy.ops.wm.save_as_mainfile(filepath=blend_path)