From 47397444b30b0f461b1633297a68300179586fda Mon Sep 17 00:00:00 2001 From: skal Date: Tue, 10 Feb 2026 08:01:25 +0100 Subject: feat: Add CNN post-processing effect with modular WGSL architecture MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements multi-layer convolutional neural network shader for stylized post-processing of 3D rendered scenes: **Core Components:** - CNNEffect: C++ effect class with single-layer rendering (expandable to multi-pass) - Modular WGSL snippets: cnn_activation, cnn_conv3x3/5x5/7x7, cnn_weights_generated - Placeholder identity-like weights for initial testing (to be replaced by trained weights) **Architecture:** - Flexible kernel sizes (3×3, 5×5, 7×7) via separate snippet files - ShaderComposer integration (#include resolution) - Residual connections (input + processed output) - Supports parallel convolutions (design ready, single conv implemented) **Size Impact:** - ~3-4 KB shader code (snippets + main shader) - ~2-4 KB weights (depends on network architecture when trained) - Total: ~5-8 KB (acceptable for 64k demo) **Testing:** - CNNEffect added to test_demo_effects.cc - 36/36 tests passing (100%) **Next Steps:** - Training script (scripts/train_cnn.py) to generate real weights - Multi-layer rendering with ping-pong textures - Weight quantization for size optimization handoff(Claude): CNN effect foundation complete, ready for training integration --- workspaces/main/assets.txt | 6 +++ workspaces/main/shaders/cnn/cnn_activation.wgsl | 18 +++++++++ workspaces/main/shaders/cnn/cnn_conv3x3.wgsl | 26 ++++++++++++ workspaces/main/shaders/cnn/cnn_conv5x5.wgsl | 26 ++++++++++++ workspaces/main/shaders/cnn/cnn_conv7x7.wgsl | 26 ++++++++++++ workspaces/main/shaders/cnn/cnn_layer.wgsl | 46 ++++++++++++++++++++++ .../main/shaders/cnn/cnn_weights_generated.wgsl | 17 ++++++++ 7 files changed, 165 insertions(+) create mode 100644 workspaces/main/shaders/cnn/cnn_activation.wgsl create mode 100644 workspaces/main/shaders/cnn/cnn_conv3x3.wgsl create mode 100644 workspaces/main/shaders/cnn/cnn_conv5x5.wgsl create mode 100644 workspaces/main/shaders/cnn/cnn_conv7x7.wgsl create mode 100644 workspaces/main/shaders/cnn/cnn_layer.wgsl create mode 100644 workspaces/main/shaders/cnn/cnn_weights_generated.wgsl (limited to 'workspaces') diff --git a/workspaces/main/assets.txt b/workspaces/main/assets.txt index ca77e21..53c8b3e 100644 --- a/workspaces/main/assets.txt +++ b/workspaces/main/assets.txt @@ -36,6 +36,12 @@ SHADER_PASSTHROUGH, NONE, shaders/passthrough.wgsl, "Passthrough Shader" SHADER_ELLIPSE, NONE, shaders/ellipse.wgsl, "Ellipse Shader" SHADER_PARTICLE_SPRAY_COMPUTE, NONE, shaders/particle_spray_compute.wgsl, "Particle Spray Compute" SHADER_GAUSSIAN_BLUR, NONE, shaders/gaussian_blur.wgsl, "Gaussian Blur Shader" +SHADER_CNN_ACTIVATION, NONE, shaders/cnn/cnn_activation.wgsl, "CNN Activation Functions" +SHADER_CNN_CONV3X3, NONE, shaders/cnn/cnn_conv3x3.wgsl, "CNN 3x3 Convolution" +SHADER_CNN_CONV5X5, NONE, shaders/cnn/cnn_conv5x5.wgsl, "CNN 5x5 Convolution" +SHADER_CNN_CONV7X7, NONE, shaders/cnn/cnn_conv7x7.wgsl, "CNN 7x7 Convolution" +SHADER_CNN_WEIGHTS, NONE, shaders/cnn/cnn_weights_generated.wgsl, "CNN Weights (Generated)" +SHADER_CNN_LAYER, NONE, shaders/cnn/cnn_layer.wgsl, "CNN Layer Shader" SHADER_SOLARIZE, NONE, shaders/solarize.wgsl, "Solarize Shader" SHADER_DISTORT, NONE, shaders/distort.wgsl, "Distort Shader" SHADER_CHROMA_ABERRATION, NONE, shaders/chroma_aberration.wgsl, "Chroma Aberration Shader" diff --git a/workspaces/main/shaders/cnn/cnn_activation.wgsl b/workspaces/main/shaders/cnn/cnn_activation.wgsl new file mode 100644 index 0000000..4fe771e --- /dev/null +++ b/workspaces/main/shaders/cnn/cnn_activation.wgsl @@ -0,0 +1,18 @@ +// CNN activation functions +// 4 functions: tanh, ReLU, sigmoid, leaky_relu + +fn cnn_tanh(x: vec4) -> vec4 { + return tanh(x); +} + +fn cnn_relu(x: vec4) -> vec4 { + return max(vec4(0.0), x); +} + +fn cnn_sigmoid(x: vec4) -> vec4 { + return 1.0 / (1.0 + exp(-x)); +} + +fn cnn_leaky_relu(x: vec4, alpha: f32) -> vec4 { + return max(alpha * x, x); +} diff --git a/workspaces/main/shaders/cnn/cnn_conv3x3.wgsl b/workspaces/main/shaders/cnn/cnn_conv3x3.wgsl new file mode 100644 index 0000000..06ca73a --- /dev/null +++ b/workspaces/main/shaders/cnn/cnn_conv3x3.wgsl @@ -0,0 +1,26 @@ +// 3x3 convolution with weight indexing +// Samples 9 pixels, applies mat4 weights per sample + +fn cnn_conv3x3( + tex: texture_2d, + samp: sampler, + uv: vec2, + resolution: vec2, + weights: array, 9>, + bias: vec4 +) -> vec4 { + let step = 1.0 / resolution; + var sum = bias; + var idx = 0; + + for (var dy = -1; dy <= 1; dy++) { + for (var dx = -1; dx <= 1; dx++) { + let offset = vec2(f32(dx), f32(dy)) * step; + let sample = textureSample(tex, samp, uv + offset); + sum += weights[idx] * sample; + idx++; + } + } + + return sum; +} diff --git a/workspaces/main/shaders/cnn/cnn_conv5x5.wgsl b/workspaces/main/shaders/cnn/cnn_conv5x5.wgsl new file mode 100644 index 0000000..3d4a03a --- /dev/null +++ b/workspaces/main/shaders/cnn/cnn_conv5x5.wgsl @@ -0,0 +1,26 @@ +// 5x5 convolution with 25 samples +// Applies mat4 weights per sample + +fn cnn_conv5x5( + tex: texture_2d, + samp: sampler, + uv: vec2, + resolution: vec2, + weights: array, 25>, + bias: vec4 +) -> vec4 { + let step = 1.0 / resolution; + var sum = bias; + var idx = 0; + + for (var dy = -2; dy <= 2; dy++) { + for (var dx = -2; dx <= 2; dx++) { + let offset = vec2(f32(dx), f32(dy)) * step; + let sample = textureSample(tex, samp, uv + offset); + sum += weights[idx] * sample; + idx++; + } + } + + return sum; +} diff --git a/workspaces/main/shaders/cnn/cnn_conv7x7.wgsl b/workspaces/main/shaders/cnn/cnn_conv7x7.wgsl new file mode 100644 index 0000000..ba28d64 --- /dev/null +++ b/workspaces/main/shaders/cnn/cnn_conv7x7.wgsl @@ -0,0 +1,26 @@ +// 7x7 convolution with 49 samples +// Applies mat4 weights per sample + +fn cnn_conv7x7( + tex: texture_2d, + samp: sampler, + uv: vec2, + resolution: vec2, + weights: array, 49>, + bias: vec4 +) -> vec4 { + let step = 1.0 / resolution; + var sum = bias; + var idx = 0; + + for (var dy = -3; dy <= 3; dy++) { + for (var dx = -3; dx <= 3; dx++) { + let offset = vec2(f32(dx), f32(dy)) * step; + let sample = textureSample(tex, samp, uv + offset); + sum += weights[idx] * sample; + idx++; + } + } + + return sum; +} diff --git a/workspaces/main/shaders/cnn/cnn_layer.wgsl b/workspaces/main/shaders/cnn/cnn_layer.wgsl new file mode 100644 index 0000000..e026ce8 --- /dev/null +++ b/workspaces/main/shaders/cnn/cnn_layer.wgsl @@ -0,0 +1,46 @@ +// CNN layer shader - uses modular convolution snippets +// Supports multi-pass rendering with residual connections + +@group(0) @binding(0) var smplr: sampler; +@group(0) @binding(1) var txt: texture_2d; + +#include "common_uniforms" +#include "cnn_activation" +#include "cnn_conv3x3" +#include "cnn_weights_generated" + +struct CNNLayerParams { + layer_index: i32, + use_residual: i32, + _pad: vec2, +}; + +@group(0) @binding(2) var uniforms: CommonUniforms; +@group(0) @binding(3) var params: CNNLayerParams; + +@vertex fn vs_main(@builtin(vertex_index) i: u32) -> @builtin(position) vec4 { + var pos = array, 3>( + vec2(-1.0, -1.0), vec2(3.0, -1.0), vec2(-1.0, 3.0) + ); + return vec4(pos[i], 0.0, 1.0); +} + +@fragment fn fs_main(@builtin(position) p: vec4) -> @location(0) vec4 { + let uv = p.xy / uniforms.resolution; + var result = vec4(0.0); + + // Single layer for now (layer 0) + if (params.layer_index == 0) { + result = cnn_conv3x3(txt, smplr, uv, uniforms.resolution, + weights_layer0, bias_layer0); + result = cnn_tanh(result); + } + + // Residual connection + if (params.use_residual != 0) { + let input = textureSample(txt, smplr, uv); + result = input + result * 0.3; + } + + return result; +} diff --git a/workspaces/main/shaders/cnn/cnn_weights_generated.wgsl b/workspaces/main/shaders/cnn/cnn_weights_generated.wgsl new file mode 100644 index 0000000..98c17ff --- /dev/null +++ b/workspaces/main/shaders/cnn/cnn_weights_generated.wgsl @@ -0,0 +1,17 @@ +// Generated CNN weights and biases +// DO NOT EDIT MANUALLY - regenerate with scripts/train_cnn.py + +// Placeholder identity-like weights for initial testing +// Layer 0: 3x3 convolution +const weights_layer0: array, 9> = array( + mat4x4(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + mat4x4(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + mat4x4(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + mat4x4(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + mat4x4(1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0), + mat4x4(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + mat4x4(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + mat4x4(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + mat4x4(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0) +); +const bias_layer0 = vec4(0.0, 0.0, 0.0, 0.0); -- cgit v1.2.3