I’ve been attempting for quite some time to construct my own screen spaced ambient occlusion, however I’ve seem to hit a roadblock when it comes to acquiring the objects’ shaded normals via color data from the framebuffer. For some reason the framebuffers color output is non-existent, my speculations lies in the rendering order confliction since I’m also attempting to render the result back to screen or it could just be that the transfer between color data and GPUTexture is wrong, but that’s just my speculation.
When display mode is initiated, all objects in the scene are assigned a material that display object’s normals in object space.
The color data read from the framebuffer is passed to a GPUTexture which is then passed to a custom SSAO shader, but like I said, the output is empty.
I’m curious to know what’s causing this? The depth data is at least correct, which is something I’ve already tested for independently.
import bpy
import gpu
import numpy as np
from gpu.types import Buffer
from gpu_extras.batch import batch_for_shader
from gpu_extras.presets import draw_texture_2d
import random
import struct
# Parameters
ssao_radius = 1.0
ssao_samples = 16
ssao_scale = 1.0
overlay_opacity = 1.0
depth_threshold = 0.99
# Global variables
_overlay_handle = None
_shader = None
_original_materials = {}
_original_viewport_shading = None
_normal_material = None
_current_size = (0, 0)
def create_normal_material():
"""Create a material that displays world space normals"""
if "WorldNormalMaterial" in bpy.data.materials:
return bpy.data.materials["WorldNormalMaterial"]
mat = bpy.data.materials.new(name="WorldNormalMaterial")
mat.use_nodes = True
nodes = mat.node_tree.nodes
links = mat.node_tree.links
# Clear default nodes
nodes.clear()
# Create nodes
geometry_node = nodes.new(type='ShaderNodeNewGeometry')
emission_node = nodes.new(type='ShaderNodeEmission')
output_node = nodes.new(type='ShaderNodeOutputMaterial')
# Position nodes
geometry_node.location = (-400, 0)
emission_node.location = (0, 0)
output_node.location = (200, 0)
# Connect nodes
links.new(geometry_node.outputs['Normal'], emission_node.inputs['Color'])
links.new(emission_node.outputs['Emission'], output_node.inputs['Surface'])
return mat
def setup_normal_viewport():
"""Setup viewport for normal rendering"""
global _original_materials, _original_viewport_shading, _normal_material
# Store original viewport shading
for area in bpy.context.screen.areas:
if area.type == 'VIEW_3D':
space = area.spaces.active
_original_viewport_shading = space.shading.type
space.shading.type = 'MATERIAL'
break
# Create normal material
_normal_material = create_normal_material()
# Store and replace materials
_original_materials.clear()
for obj in bpy.context.scene.objects:
if hasattr(obj, 'data') and hasattr(obj.data, 'materials'):
_original_materials[obj] = list(obj.data.materials) if obj.data.materials else []
obj.data.materials.clear()
if obj.type in {'MESH', 'CURVE', 'SURFACE', 'FONT', 'META'}:
obj.data.materials.append(_normal_material)
def restore_original_viewport():
"""Restore original viewport state"""
global _original_materials, _original_viewport_shading
# Restore materials
for obj, materials in _original_materials.items():
if hasattr(obj, 'data') and hasattr(obj.data, 'materials'):
obj.data.materials.clear()
for mat in materials:
obj.data.materials.append(mat)
_original_materials.clear()
# Restore viewport shading
if _original_viewport_shading:
for area in bpy.context.screen.areas:
if area.type == 'VIEW_3D':
space = area.spaces.active
space.shading.type = _original_viewport_shading
break
def create_ssao_shader():
"""Create SSAO shader"""
vertex_shader = '''
in vec2 pos;
in vec2 texCoord;
out vec2 texCoord_interp;
void main()
{
gl_Position = vec4(pos, 0.0, 1.0);
texCoord_interp = texCoord;
}
'''
fragment_shader = '''
in vec2 texCoord_interp;
out vec4 fragColor;
uniform sampler2D depth_texture;
uniform sampler2D normal_texture;
uniform float radius;
uniform int sampleCount;
uniform float scale;
uniform float threshold;
uniform float opacity;
uniform float near;
uniform float far;
uniform mat4 inv_proj_matrix;
uniform mat4 inv_view_matrix;
uniform mat4 proj_matrix;
uniform mat4 view_matrix;
// Sample kernel - we'll use a fixed maximum size
const int MAX_SAMPLES = 64;
uniform vec3 samples[MAX_SAMPLES];
vec3 world_pos_from_depth(float depth) {
float z = depth * 2.0 - 1.0;
vec4 clip_space_position = vec4(texCoord_interp * 2.0 - 1.0, z, 1.0);
vec4 view_space_position = inv_proj_matrix * clip_space_position;
// Perspective division
view_space_position /= view_space_position.w;
vec4 world_space_position = inv_view_matrix * view_space_position;
return world_space_position.xyz;
}
float nonlinear_to_linear(float depth)
{
return (2.0*near) / (far + near - depth*(far - near));
}
void main()
{
float depth = texture(depth_texture, texCoord_interp).r;
float linear_depth = nonlinear_to_linear(depth);
if (linear_depth >= threshold) {
discard;
}
vec3 world_pos = world_pos_from_depth(depth);
vec3 normal = texture(normal_texture, texCoord_interp).xyz;
float occlusion = 0.0;
int numSamples = min(sampleCount, MAX_SAMPLES);
for (int i = 0; i < numSamples; i++) {
// Get sample position in world space
vec3 sample_offset = samples[i] * radius;
vec3 sample_pos = world_pos + sample_offset;
// Transform sample position to view space then clip space
vec4 sample_view_pos = view_matrix * vec4(sample_pos, 1.0);
vec4 sample_clip_pos = proj_matrix * sample_view_pos;
// Perspective division
sample_clip_pos.xyz /= sample_clip_pos.w;
// Convert to UV coordinates
vec2 sample_uv = sample_clip_pos.xy * 0.5 + 0.5;
//sample_uv.y = 1.0 - sample_uv.y; // Flip Y
// Check if sample is within screen bounds
if (sample_uv.x < 0.0 || sample_uv.x > 1.0 || sample_uv.y < 0.0 || sample_uv.y > 1.0) {
continue;
}
// Sample depth at offset position
float sample_depth = texture(depth_texture, sample_uv).r;
if (sample_depth >= threshold) continue;
vec3 sample_world_pos = world_pos_from_depth(sample_depth);
// Calculate occlusion
vec3 sample_dir = normalize(sample_world_pos - world_pos);
float sample_dist = length(sample_world_pos - world_pos);
// Range check and occlusion accumulation
float range_check = smoothstep(0.0, 1.0, radius / abs(sample_dist));
float NdotS = max(dot(normal, sample_dir), 0.0);
occlusion += range_check * NdotS;
}
occlusion = 1.0 - (occlusion / float(numSamples)) * scale;
fragColor = vec4(vec3(occlusion), opacity);
}
'''
return gpu.types.GPUShader(vertex_shader, fragment_shader)
def generate_samples(num_samples):
"""Generate sample kernel for SSAO"""
samples = []
for i in range(num_samples):
# Generate sample in hemisphere oriented along Z-axis
sample = np.array([
random.uniform(-1, 1),
random.uniform(-1, 1),
random.uniform(0, 1) # Hemisphere (Z > 0)
])
sample = sample / np.linalg.norm(sample)
# Scale for distribution (more samples closer to origin)
scale = float(i) / float(num_samples)
scale = 0.1 + 0.9 * scale * scale
sample *= scale
samples.append(sample)
while len(samples) < 64:
samples.append(np.array([0.0, 0.0, 0.0]))
return samples
def get_buffer_data():
"""Extract buffer data (depth or color)"""
try:
framebuffer = gpu.state.active_framebuffer_get()
viewport = gpu.state.viewport_get()
width, height = viewport[2], viewport[3]
if width == 0 or height == 0:
return None, 0, 0
buffer_size = width * height
depth_buffer = Buffer('FLOAT', buffer_size)
framebuffer.read_depth(0, 0, width, height, data=depth_buffer)
print("Getting color")
buffer_size = width * height * 4
normal_buffer = Buffer('FLOAT', buffer_size)
framebuffer.read_color(0, 0, width, height, 4, 0, 'FLOAT', data=normal_buffer)
return normal_buffer, depth_buffer, width, height
except Exception as e:
print(f"Error getting {buffer_type} data: {e}")
return None, None, 0, 0
def draw_ssao_overlay():
"""Draw callback that renders SSAO"""
global _shader
scene = bpy.context.scene
overlay_opacity = scene.overlay_opacity
ssao_radius = scene.ssao_radius
ssao_samples = scene.ssao_samples
ssao_scale = scene.ssao_scale
# Get depth and normal data
normal_buffer, depth_buffer, width, height = get_buffer_data()
if depth_buffer is None or normal_buffer is None:
return
# Create textures
depth_texture = gpu.types.GPUTexture((width, height), format='R32F', data=depth_buffer)
normal_texture = gpu.types.GPUTexture((width, height), format='RGBA32F', data=normal_buffer)
# Create shader if it doesn't exist
if _shader is None:
_shader = create_ssao_shader()
if _shader is None:
print("Failed to create SSAO shader")
return
# Get camera matrices from 3D view
for area in bpy.context.screen.areas:
if area.type == 'VIEW_3D':
space = area.spaces.active
region_3d = space.region_3d
far = space.clip_end
near = space.clip_start
break
# Get matrices
proj_matrix = gpu.matrix.get_projection_matrix()
inv_proj_matrix = proj_matrix.inverted()
view_matrix = gpu.matrix.get_model_view_matrix()
inv_view_matrix = view_matrix.inverted()
# Generate samples
samples = generate_samples(ssao_samples)
batch = batch_for_shader(
_shader, 'TRI_FAN',
{
"pos": ((-1.0, -1.0), (1.0, -1.0), (1.0, 1.0), (-1.0, 1.0)),
"texCoord": ((0, 0), (1, 0), (1, 1), (0, 1)),
},
)
# Draw SSAO
gpu.state.depth_test_set('NONE')
gpu.state.blend_set('ALPHA')
_shader.bind()
# Set textures
_shader.uniform_sampler("normal_texture", normal_texture)
_shader.uniform_sampler("depth_texture", depth_texture)
# Set scalar uniforms
_shader.uniform_float("threshold", depth_threshold)
_shader.uniform_float("opacity", overlay_opacity)
_shader.uniform_float("near", near)
_shader.uniform_float("far", far)
_shader.uniform_float("radius", ssao_radius)
_shader.uniform_int("sampleCount", ssao_samples)
_shader.uniform_float("scale", ssao_scale)
# Set matrix uniforms
inv_proj_matrix_flat = [element for row in inv_proj_matrix.transposed() for element in row]
inv_proj_matrix_bytes = struct.pack('<16f', *inv_proj_matrix_flat)
inv_proj_matrix_loc = _shader.uniform_from_name("inv_proj_matrix")
_shader.uniform_vector_float(inv_proj_matrix_loc, inv_proj_matrix_bytes, 16)
inv_view_matrix_flat = [element for row in inv_view_matrix.transposed() for element in row]
inv_view_matrix_bytes = struct.pack('<16f', *inv_view_matrix_flat)
inv_view_matrix_loc = _shader.uniform_from_name("inv_view_matrix")
_shader.uniform_vector_float(inv_view_matrix_loc, inv_view_matrix_bytes, 16)
proj_matrix_flat = [element for row in proj_matrix.transposed() for element in row]
proj_matrix_bytes = struct.pack('<16f', *proj_matrix_flat)
proj_matrix_loc = _shader.uniform_from_name("proj_matrix")
_shader.uniform_vector_float(proj_matrix_loc, proj_matrix_bytes, 16)
view_matrix_flat = [element for row in view_matrix.transposed() for element in row]
view_matrix_bytes = struct.pack('<16f', *view_matrix_flat)
view_matrix_loc = _shader.uniform_from_name("view_matrix")
_shader.uniform_vector_float(view_matrix_loc, view_matrix_bytes, 16)
# Set samples array
samples_flat = []
for sample in samples:
samples_flat.extend(sample)
samples_bytes = struct.pack('<{}f'.format(len(samples_flat)), *samples_flat)
samples_loc = _shader.uniform_from_name("samples")
_shader.uniform_vector_float(samples_loc, samples_bytes, len(samples_flat))
batch.draw(_shader)
gpu.state.depth_mask_set(False)
draw_texture_2d(normal_texture, (0, 0), 512, 512)
def start_ssao():
"""Start SSAO overlay"""
global _overlay_handle
if _overlay_handle is None:
setup_normal_viewport()
_overlay_handle = bpy.types.SpaceView3D.draw_handler_add(draw_ssao_overlay, (), 'WINDOW', 'POST_PIXEL')
print("SSAO started")
return True
return False
def stop_ssao():
"""Stop SSAO overlay"""
global _overlay_handle
if _overlay_handle is not None:
bpy.types.SpaceView3D.draw_handler_remove(_overlay_handle, 'WINDOW')
_overlay_handle = None
restore_original_viewport()
print("SSAO stopped")
return True
return False
# Operator and Panel
class SSAOToggleOperator(bpy.types.Operator):
bl_idname = "view3d.ssao_toggle"
bl_label = "Toggle SSAO"
bl_description = "Toggle Screen Space Ambient Occlusion"
def execute(self, context):
if _overlay_handle is None:
if start_ssao():
self.report({'INFO'}, "SSAO started")
else:
self.report({'ERROR'}, "Failed to start SSAO")
else:
stop_ssao()
self.report({'INFO'}, "SSAO stopped")
for window in context.window_manager.windows:
for area in window.screen.areas:
if area.type == 'VIEW_3D':
area.tag_redraw()
return {'FINISHED'}
class SSAOPanel(bpy.types.Panel):
bl_label = "SSAO Overlay"
bl_idname = "VIEW3D_PT_ssao_overlay"
bl_space_type = 'VIEW_3D'
bl_region_type = 'UI'
bl_category = "Tool"
def draw(self, context):
layout = self.layout
if _overlay_handle is None:
layout.operator("view3d.ssao_toggle", text="Start SSAO", icon='PLAY')
else:
layout.operator("view3d.ssao_toggle", text="Stop SSAO", icon='PAUSE')
layout.separator()
layout.prop(context.scene, "ssao_radius", slider=True)
layout.prop(context.scene, "ssao_samples")
layout.prop(context.scene, "ssao_scale", slider=True)
layout.prop(context.scene, "overlay_opacity", slider=True)
# Property definitions
def register_properties():
bpy.types.Scene.ssao_radius = bpy.props.FloatProperty(
name="Radius",
description="SSAO sampling radius",
default=1.0,
min=0.001,
max=10.0
)
bpy.types.Scene.ssao_samples = bpy.props.IntProperty(
name="Samples",
description="Number of SSAO samples",
default=16,
min=1,
max=64
)
bpy.types.Scene.ssao_scale = bpy.props.FloatProperty(
name="Scale",
description="SSAO effect scale",
default=1.0,
min=0.0,
max=64.0
)
bpy.types.Scene.overlay_opacity = bpy.props.FloatProperty(
name="Opacity",
description="Overlay opacity",
default=1.0,
min=0.0,
max=1.0
)
def unregister_properties():
del bpy.types.Scene.ssao_radius
del bpy.types.Scene.ssao_samples
del bpy.types.Scene.ssao_scale
del bpy.types.Scene.overlay_opacity
def register():
register_properties()
bpy.utils.register_class(SSAOToggleOperator)
bpy.utils.register_class(SSAOPanel)
def unregister():
stop_ssao()
bpy.utils.unregister_class(SSAOToggleOperator)
bpy.utils.unregister_class(SSAOPanel)
unregister_properties()
if __name__ == "__main__":
register()