format
This commit is contained in:
hofee
parent
950b1473b2
commit
5905669dbd
161
blender_util.py
161
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"
|
||||
@ -35,14 +35,11 @@ class BlenderUtils:
|
||||
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)
|
||||
@ -50,13 +47,15 @@ class BlenderUtils:
|
||||
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)
|
||||
@ -67,11 +66,11 @@ 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 # 关闭辉光
|
||||
@ -80,20 +79,29 @@ class BlenderUtils:
|
||||
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.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_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)
|
||||
|
||||
@ -101,9 +109,9 @@ class BlenderUtils:
|
||||
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
|
||||
@ -139,16 +147,24 @@ class BlenderUtils:
|
||||
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
|
||||
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):
|
||||
@ -179,7 +195,9 @@ class BlenderUtils:
|
||||
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)
|
||||
@ -191,13 +209,14 @@ class BlenderUtils:
|
||||
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,71 +228,69 @@ 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 = 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
|
||||
|
||||
@ -319,7 +336,14 @@ class BlenderUtils:
|
||||
@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)
|
||||
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user