optimize code structure
This commit is contained in:
hofee
@@ -6,8 +6,9 @@ import bpy
|
||||
import numpy as np
|
||||
import mathutils
|
||||
import requests
|
||||
from blender.blender_util import BlenderUtils
|
||||
from blender.view_sample_util import ViewSampleUtil
|
||||
from utils.blender_util import BlenderUtils
|
||||
from utils.view_sample_util import ViewSampleUtil
|
||||
from utils.material_util import MaterialUtil
|
||||
|
||||
class DataGenerator:
|
||||
def __init__(self, config):
|
||||
@@ -103,29 +104,7 @@ class DataGenerator:
|
||||
bpy.context.object.rigid_body.type = 'PASSIVE'
|
||||
bpy.ops.object.shade_auto_smooth()
|
||||
|
||||
# 创建不受光照影响的材质
|
||||
mat = bpy.data.materials.new(name="RedMaterial")
|
||||
mat.use_nodes = True
|
||||
|
||||
# 清除默认节点
|
||||
nodes = mat.node_tree.nodes
|
||||
for node in nodes:
|
||||
nodes.remove(node)
|
||||
|
||||
# 添加 Emission 节点
|
||||
emission_node = nodes.new(type='ShaderNodeEmission')
|
||||
emission_node.inputs['Color'].default_value = (1.0, 0.0, 0.0, 1.0) # 红色
|
||||
|
||||
# 添加 Material Output 节点
|
||||
output_node = nodes.new(type='ShaderNodeOutputMaterial')
|
||||
|
||||
# 连接节点
|
||||
links = mat.node_tree.links
|
||||
links.new(emission_node.outputs['Emission'], output_node.inputs['Surface'])
|
||||
|
||||
# 将材质赋给对象
|
||||
platform.data.materials.clear()
|
||||
platform.data.materials.append(mat)
|
||||
MaterialUtil.change_object_material(platform, MaterialUtil.create_mask_material(color=(1.0, 0, 0)))
|
||||
|
||||
self.display_table_config = {
|
||||
"height": height,
|
||||
@@ -166,31 +145,7 @@ class DataGenerator:
|
||||
|
||||
bpy.ops.rigidbody.object_add()
|
||||
bpy.context.object.rigid_body.type = 'ACTIVE'
|
||||
|
||||
# 创建不受光照影响的材质
|
||||
mat = bpy.data.materials.new(name="GreenMaterial")
|
||||
mat.use_nodes = True
|
||||
|
||||
# 清除默认节点
|
||||
nodes = mat.node_tree.nodes
|
||||
for node in nodes:
|
||||
nodes.remove(node)
|
||||
|
||||
# 添加 Emission 节点
|
||||
emission_node = nodes.new(type='ShaderNodeEmission')
|
||||
emission_node.inputs['Color'].default_value = (0.0, 1.0, 0.0, 1.0) # 绿色
|
||||
|
||||
# 添加 Material Output 节点
|
||||
output_node = nodes.new(type='ShaderNodeOutputMaterial')
|
||||
|
||||
# 连接节点
|
||||
links = mat.node_tree.links
|
||||
links.new(emission_node.outputs['Emission'], output_node.inputs['Surface'])
|
||||
|
||||
# 将材质赋给对象
|
||||
obj.data.materials.clear()
|
||||
obj.data.materials.append(mat)
|
||||
|
||||
MaterialUtil.change_object_material(obj, MaterialUtil.create_mask_material(color=(0, 1.0, 0)))
|
||||
self.target_obj = obj
|
||||
|
||||
|
||||
@@ -249,9 +204,7 @@ class DataGenerator:
|
||||
self.set_progress("render frame", len(view_data["cam_poses"]), len(view_data["cam_poses"]))
|
||||
BlenderUtils.save_scene_info(scene_dir, self.display_table_config, object_name)
|
||||
|
||||
|
||||
|
||||
self.change_target_obj_material_to_normal()
|
||||
MaterialUtil.change_object_material(self.target_obj, MaterialUtil.create_normal_material())
|
||||
for i, cam_pose in enumerate(view_data["cam_poses"]):
|
||||
BlenderUtils.set_camera_at(cam_pose)
|
||||
BlenderUtils.render_normal_and_depth(scene_dir, f"{i}", binocular_vision=self.binocular_vision, target_object = self.target_obj)
|
||||
@@ -269,196 +222,6 @@ class DataGenerator:
|
||||
|
||||
return True
|
||||
|
||||
def change_target_obj_material_to_normal(self):
|
||||
|
||||
material_name = "normal"
|
||||
mat = bpy.data.materials.get(material_name) or bpy.data.materials.new(
|
||||
material_name
|
||||
)
|
||||
mat.use_nodes = True
|
||||
node_tree = mat.node_tree
|
||||
nodes = node_tree.nodes
|
||||
nodes.clear()
|
||||
|
||||
links = node_tree.links
|
||||
links.clear()
|
||||
|
||||
# Nodes:
|
||||
new_node = nodes.new(type="ShaderNodeMath")
|
||||
# new_node.active_preview = False
|
||||
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
|
||||
new_node.location = (151.59744262695312, 854.5482177734375)
|
||||
new_node.name = "Math"
|
||||
new_node.operation = "MULTIPLY"
|
||||
new_node.select = False
|
||||
new_node.use_clamp = False
|
||||
new_node.width = 140.0
|
||||
new_node.inputs[0].default_value = 0.5
|
||||
new_node.inputs[1].default_value = 1.0
|
||||
new_node.inputs[2].default_value = 0.0
|
||||
new_node.outputs[0].default_value = 0.0
|
||||
|
||||
new_node = nodes.new(type="ShaderNodeLightPath")
|
||||
# new_node.active_preview = False
|
||||
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
|
||||
new_node.location = (602.9912719726562, 1046.660888671875)
|
||||
new_node.name = "Light Path"
|
||||
new_node.select = False
|
||||
new_node.width = 140.0
|
||||
new_node.outputs[0].default_value = 0.0
|
||||
new_node.outputs[1].default_value = 0.0
|
||||
new_node.outputs[2].default_value = 0.0
|
||||
new_node.outputs[3].default_value = 0.0
|
||||
new_node.outputs[4].default_value = 0.0
|
||||
new_node.outputs[5].default_value = 0.0
|
||||
new_node.outputs[6].default_value = 0.0
|
||||
new_node.outputs[7].default_value = 0.0
|
||||
new_node.outputs[8].default_value = 0.0
|
||||
new_node.outputs[9].default_value = 0.0
|
||||
new_node.outputs[10].default_value = 0.0
|
||||
new_node.outputs[11].default_value = 0.0
|
||||
new_node.outputs[12].default_value = 0.0
|
||||
|
||||
new_node = nodes.new(type="ShaderNodeOutputMaterial")
|
||||
# new_node.active_preview = False
|
||||
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
|
||||
new_node.is_active_output = True
|
||||
new_node.location = (1168.93017578125, 701.84033203125)
|
||||
new_node.name = "Material Output"
|
||||
new_node.select = False
|
||||
new_node.target = "ALL"
|
||||
new_node.width = 140.0
|
||||
new_node.inputs[2].default_value = [0.0, 0.0, 0.0]
|
||||
|
||||
new_node = nodes.new(type="ShaderNodeBsdfTransparent")
|
||||
# new_node.active_preview = False
|
||||
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
|
||||
new_node.location = (731.72900390625, 721.4832763671875)
|
||||
new_node.name = "Transparent BSDF"
|
||||
new_node.select = False
|
||||
new_node.width = 140.0
|
||||
new_node.inputs[0].default_value = [1.0, 1.0, 1.0, 1.0]
|
||||
|
||||
new_node = nodes.new(type="ShaderNodeCombineXYZ")
|
||||
# new_node.active_preview = False
|
||||
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
|
||||
new_node.location = (594.4229736328125, 602.9271240234375)
|
||||
new_node.name = "Combine XYZ"
|
||||
new_node.select = False
|
||||
new_node.width = 140.0
|
||||
new_node.inputs[0].default_value = 0.0
|
||||
new_node.inputs[1].default_value = 0.0
|
||||
new_node.inputs[2].default_value = 0.0
|
||||
new_node.outputs[0].default_value = [0.0, 0.0, 0.0]
|
||||
|
||||
new_node = nodes.new(type="ShaderNodeMixShader")
|
||||
# new_node.active_preview = False
|
||||
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
|
||||
new_node.location = (992.7239990234375, 707.2142333984375)
|
||||
new_node.name = "Mix Shader"
|
||||
new_node.select = False
|
||||
new_node.width = 140.0
|
||||
new_node.inputs[0].default_value = 0.5
|
||||
|
||||
new_node = nodes.new(type="ShaderNodeEmission")
|
||||
# new_node.active_preview = False
|
||||
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
|
||||
new_node.location = (774.0802612304688, 608.2547607421875)
|
||||
new_node.name = "Emission"
|
||||
new_node.select = False
|
||||
new_node.width = 140.0
|
||||
new_node.inputs[0].default_value = [1.0, 1.0, 1.0, 1.0]
|
||||
new_node.inputs[1].default_value = 1.0
|
||||
|
||||
new_node = nodes.new(type="ShaderNodeSeparateXYZ")
|
||||
# new_node.active_preview = False
|
||||
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
|
||||
new_node.location = (-130.12167358398438, 558.1497802734375)
|
||||
new_node.name = "Separate XYZ"
|
||||
new_node.select = False
|
||||
new_node.width = 140.0
|
||||
new_node.inputs[0].default_value = [0.0, 0.0, 0.0]
|
||||
new_node.outputs[0].default_value = 0.0
|
||||
new_node.outputs[1].default_value = 0.0
|
||||
new_node.outputs[2].default_value = 0.0
|
||||
|
||||
new_node = nodes.new(type="ShaderNodeMath")
|
||||
# new_node.active_preview = False
|
||||
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
|
||||
new_node.location = (162.43240356445312, 618.8094482421875)
|
||||
new_node.name = "Math.002"
|
||||
new_node.operation = "MULTIPLY"
|
||||
new_node.select = False
|
||||
new_node.use_clamp = False
|
||||
new_node.width = 140.0
|
||||
new_node.inputs[0].default_value = 0.5
|
||||
new_node.inputs[1].default_value = 1.0
|
||||
new_node.inputs[2].default_value = 0.0
|
||||
new_node.outputs[0].default_value = 0.0
|
||||
|
||||
new_node = nodes.new(type="ShaderNodeMath")
|
||||
# new_node.active_preview = False
|
||||
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
|
||||
new_node.location = (126.8158187866211, 364.5539855957031)
|
||||
new_node.name = "Math.001"
|
||||
new_node.operation = "MULTIPLY"
|
||||
new_node.select = False
|
||||
new_node.use_clamp = False
|
||||
new_node.width = 140.0
|
||||
new_node.inputs[0].default_value = 0.5
|
||||
new_node.inputs[1].default_value = -1.0
|
||||
new_node.inputs[2].default_value = 0.0
|
||||
new_node.outputs[0].default_value = 0.0
|
||||
|
||||
new_node = nodes.new(type="ShaderNodeVectorTransform")
|
||||
# new_node.active_preview = False
|
||||
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
|
||||
new_node.convert_from = "WORLD"
|
||||
new_node.convert_to = "CAMERA"
|
||||
new_node.location = (-397.0209045410156, 594.7037353515625)
|
||||
new_node.name = "Vector Transform"
|
||||
new_node.select = False
|
||||
new_node.vector_type = "VECTOR"
|
||||
new_node.width = 140.0
|
||||
new_node.inputs[0].default_value = [0.5, 0.5, 0.5]
|
||||
new_node.outputs[0].default_value = [0.0, 0.0, 0.0]
|
||||
|
||||
new_node = nodes.new(type="ShaderNodeNewGeometry")
|
||||
# new_node.active_preview = False
|
||||
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
|
||||
new_node.location = (-651.8067016601562, 593.0455932617188)
|
||||
new_node.name = "Geometry"
|
||||
new_node.width = 140.0
|
||||
new_node.outputs[0].default_value = [0.0, 0.0, 0.0]
|
||||
new_node.outputs[1].default_value = [0.0, 0.0, 0.0]
|
||||
new_node.outputs[2].default_value = [0.0, 0.0, 0.0]
|
||||
new_node.outputs[3].default_value = [0.0, 0.0, 0.0]
|
||||
new_node.outputs[4].default_value = [0.0, 0.0, 0.0]
|
||||
new_node.outputs[5].default_value = [0.0, 0.0, 0.0]
|
||||
new_node.outputs[6].default_value = 0.0
|
||||
new_node.outputs[7].default_value = 0.0
|
||||
new_node.outputs[8].default_value = 0.0
|
||||
|
||||
# Links :
|
||||
|
||||
links.new(nodes["Light Path"].outputs[0], nodes["Mix Shader"].inputs[0])
|
||||
links.new(nodes["Separate XYZ"].outputs[0], nodes["Math"].inputs[0])
|
||||
links.new(nodes["Separate XYZ"].outputs[1], nodes["Math.002"].inputs[0])
|
||||
links.new(nodes["Separate XYZ"].outputs[2], nodes["Math.001"].inputs[0])
|
||||
links.new(nodes["Vector Transform"].outputs[0], nodes["Separate XYZ"].inputs[0])
|
||||
links.new(nodes["Combine XYZ"].outputs[0], nodes["Emission"].inputs[0])
|
||||
links.new(nodes["Math"].outputs[0], nodes["Combine XYZ"].inputs[0])
|
||||
links.new(nodes["Math.002"].outputs[0], nodes["Combine XYZ"].inputs[1])
|
||||
links.new(nodes["Math.001"].outputs[0], nodes["Combine XYZ"].inputs[2])
|
||||
links.new(nodes["Transparent BSDF"].outputs[0], nodes["Mix Shader"].inputs[1])
|
||||
links.new(nodes["Emission"].outputs[0], nodes["Mix Shader"].inputs[2])
|
||||
links.new(nodes["Mix Shader"].outputs[0], nodes["Material Output"].inputs[0])
|
||||
links.new(nodes["Geometry"].outputs[1], nodes["Vector Transform"].inputs[0])
|
||||
|
||||
self.target_obj.data.materials.clear()
|
||||
self.target_obj.data.materials.append(mat)
|
||||
|
||||
|
||||
|
||||
def simulate_scene(self, frame_limit=120, depth = 0, diag = 0):
|
||||
bpy.context.view_layer.update()
|
||||
|
||||
Reference in New Issue
Block a user