feat(object,render): add ambient occlusion effect

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])

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.