341 lines
14 KiB
Python
341 lines
14 KiB
Python
|
|
import os
|
|
import json
|
|
import bpy
|
|
import time
|
|
import gc
|
|
import numpy as np
|
|
import mathutils
|
|
|
|
class BlenderUtils:
|
|
|
|
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, BlenderUtils.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 = 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)
|
|
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 add_table(table_model_path):
|
|
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'
|
|
|
|
mat = bpy.data.materials.new(name="TableYellowMaterial")
|
|
mat.diffuse_color = (1.0, 1.0, 0.0, 1.0)
|
|
if len(table.data.materials) > 0:
|
|
table.data.materials[0] = mat
|
|
else:
|
|
table.data.materials.append(mat)
|
|
|
|
@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.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
|
|
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':
|
|
light.data.energy = config["power"]
|
|
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.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.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)
|
|
|
|
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 = BlenderUtils.get_obj(BlenderUtils.CAMERA_NAME)
|
|
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
|
|
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
|
|
|
|
@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_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)
|
|
for cam_name in target_cameras:
|
|
bpy.context.scene.camera = BlenderUtils.get_obj(cam_name)
|
|
bpy.context.scene.view_layers["ViewLayer"].use_pass_z = True
|
|
bpy.context.scene.view_layers["ViewLayer"].use_pass_normal = True
|
|
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.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])
|
|
|
|
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'
|
|
tree.links.new(rl.outputs['Normal'], 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
|
|
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 = BlenderUtils.get_obj(BlenderUtils.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 = BlenderUtils.get_obj(BlenderUtils.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(BlenderUtils.CAMERA_OBJECT_NAME)
|
|
keep_objects.add(BlenderUtils.CAMERA_NAME)
|
|
keep_objects.add(BlenderUtils.CAMERA_RIGHT_NAME)
|
|
keep_objects.add(BlenderUtils.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(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:
|
|
obj_info = {
|
|
"location": list(obj.location),
|
|
"rotation_euler": list(obj.rotation_euler),
|
|
"scale": list(obj.scale)
|
|
}
|
|
scene_info[obj.name] = obj_info
|
|
scene_info[BlenderUtils.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)
|
|
|
|
|
|
|