From 7132ff83cea118167dcbfde395994c7f5460b544 Mon Sep 17 00:00:00 2001 From: hofee <64160135+GitHofee@users.noreply.github.com> Date: Tue, 15 Oct 2024 11:17:26 -0500 Subject: [PATCH] add normal material --- blender_util.py | 77 +++++++---------- data_generator.py | 209 ++++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 233 insertions(+), 53 deletions(-) diff --git a/blender_util.py b/blender_util.py index ada8f44..8a75016 100644 --- a/blender_util.py +++ b/blender_util.py @@ -195,7 +195,38 @@ class BlenderUtils: return min_z @staticmethod - def render_and_save( + def render_normal( + 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) + 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, "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}.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 render_mask_and_depth( output_dir, file_name, binocular_vision=False, target_object=None ): target_cameras = [BlenderUtils.CAMERA_NAME] @@ -203,8 +234,6 @@ class BlenderUtils: 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 = "" @@ -247,48 +276,6 @@ class BlenderUtils: 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.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") - - # 设置映射范围 - 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 - - # 分别连接到法线的 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]) - - # 合并 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"]) - - # 输出到文件 - 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(combine_rgb.outputs[0], output_normal.inputs[0]) bpy.ops.render.render(write_still=True) msg = "success" diff --git a/data_generator.py b/data_generator.py index 698fd3b..ed37170 100644 --- a/data_generator.py +++ b/data_generator.py @@ -243,7 +243,7 @@ class DataGenerator: np.savetxt(os.path.join(scene_dir, "points_and_normals.txt"), points_normals) for i, cam_pose in enumerate(view_data["cam_poses"]): BlenderUtils.set_camera_at(cam_pose) - BlenderUtils.render_and_save(scene_dir, f"{i}", binocular_vision=self.binocular_vision, target_object = self.target_obj) + BlenderUtils.render_mask_and_depth(scene_dir, f"{i}", binocular_vision=self.binocular_vision, target_object = self.target_obj) BlenderUtils.save_cam_params(scene_dir, i, binocular_vision=self.binocular_vision) self.set_progress("render frame", i, len(view_data["cam_poses"])) self.set_progress("render frame", len(view_data["cam_poses"]), len(view_data["cam_poses"])) @@ -255,15 +255,208 @@ class DataGenerator: file_path = os.path.join(depth_dir, depth_file) new_file_path = os.path.join(depth_dir, f"{name}.png") os.rename(file_path,new_file_path) - normal_dir = os.path.join(scene_dir, "normal") - for normal_file in os.listdir(normal_dir): - if not normal_file.endswith(".png"): - name, _ = os.path.splitext(normal_file) - file_path = os.path.join(normal_dir, normal_file) - new_file_path = os.path.join(normal_dir, f"{name}.png") - os.rename(file_path,new_file_path) + + self.change_target_obj_material_to_normal() + for i, cam_pose in enumerate(view_data["cam_poses"]): + BlenderUtils.set_camera_at(cam_pose) + BlenderUtils.render_normal(scene_dir, f"{i}", binocular_vision=self.binocular_vision, target_object = self.target_obj) + BlenderUtils.save_cam_params(scene_dir, i, binocular_vision=self.binocular_vision) + self.set_progress("render normal frame", i, len(view_data["cam_poses"])) + self.set_progress("render normal frame", len(view_data["cam_poses"]), len(view_data["cam_poses"])) + 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() bpy.ops.screen.animation_play()