diff --git a/blenderproc/python/renderer/RendererUtility.py b/blenderproc/python/renderer/RendererUtility.py index 5ef699b0e..3a124e5ba 100644 --- a/blenderproc/python/renderer/RendererUtility.py +++ b/blenderproc/python/renderer/RendererUtility.py @@ -918,3 +918,38 @@ def set_render_devices(use_only_cpu: bool = False, desired_gpu_device_type: Unio bpy.context.scene.cycles.device = "CPU" bpy.context.preferences.addons['cycles'].preferences.compute_device_type = "NONE" print("Using only the CPU for rendering") + + +def enable_ambient_occlusion(distance: float = 1.0, strength: float = 1.0): + """ Enables ambient occlusion in the scene. This can improve object part or point separation. + + :param distance: Length of rays, defines how far away other faces give occlusion effect. + :param strength: The amount of ambient occlusion effect added to the scene. + """ + bpy.context.scene.view_layers['ViewLayer'].use_pass_ambient_occlusion = True + bpy.context.scene.world.light_settings.distance = distance + bpy.context.scene.use_nodes = True + + nodes = bpy.context.scene.node_tree.nodes + render_layers_node = Utility.get_the_one_node_with_type(nodes, 'CompositorNodeRLayers') + mix_rgb_node = nodes.new('CompositorNodeMixRGB') + mix_rgb_node.blend_type = 'MULTIPLY' + mix_rgb_node.inputs[0].default_value = strength + + # Insert the mix-rgb node in between denoise and composite node, if a denoise node is present + links = bpy.context.scene.node_tree.links + try: + denoise_node = Utility.get_the_one_node_with_type(nodes, 'CompositorNodeDenoise') + Utility.insert_node_instead_existing_link(links, + render_layers_node.outputs['Image'], + mix_rgb_node.inputs['Image'], + mix_rgb_node.outputs['Image'], + denoise_node.inputs['Image']) + except RuntimeError: + composite_node = Utility.get_the_one_node_with_type(nodes, 'CompositorNodeComposite') + Utility.insert_node_instead_existing_link(links, + render_layers_node.outputs['Image'], + mix_rgb_node.inputs['Image'], + mix_rgb_node.outputs['Image'], + composite_node.inputs['Image']) + links.new(render_layers_node.outputs['AO'], mix_rgb_node.inputs[2]) diff --git a/blenderproc/python/types/MeshObjectUtility.py b/blenderproc/python/types/MeshObjectUtility.py index f62f70e5e..6ba3c6ac8 100644 --- a/blenderproc/python/types/MeshObjectUtility.py +++ b/blenderproc/python/types/MeshObjectUtility.py @@ -551,6 +551,31 @@ def add_auto_smooth_modifier(self, angle: float = 30.0): modifier = self.blender_obj.modifiers["Smooth by Angle"] modifier["Input_1"] = np.deg2rad(float(angle)) + def add_ambient_occlusion(self, distance: float = 1.0, samples: int = 16, only_local: bool = False): + """ Adds an ambient occlusion effect to the object. This can improve object part or point separation. + + :param distance: Length of rays, defines how far away other faces give occlusion effect. + :param samples: Number of rays to trace per shader evaluation. + :param only_local: Only consider the object itself when computing AO. + """ + for material in self.get_materials(): + ao_node = material.new_node('ShaderNodeAmbientOcclusion') + ao_node.inputs['Distance'].default_value = distance + ao_node.samples = samples + ao_node.only_local = only_local + + # Insert the ambient occlusion node in between the attribute node and the bsdf node, if present + try: + attribute_node = material.get_the_one_node_with_type('ShaderNodeAttribute') + bsdf_node = material.get_the_one_node_with_type('ShaderNodeBsdfPrincipled') + material.insert_node_instead_existing_link(attribute_node.outputs['Color'], + ao_node.inputs['Color'], + ao_node.outputs['Color'], + bsdf_node.inputs['Base Color']) + except RuntimeError: + ao_node.inputs['Color'].default_value = material.get_principled_shader_value('Base Color') + material.set_principled_shader_value('Base Color', ao_node.outputs['Color']) + def mesh_as_trimesh(self) -> Trimesh: """ Returns a trimesh.Trimesh instance of the MeshObject.