diff options
| author | Maze Winther <mazewinther@gmail.com> | 2026-04-07 01:09:13 +0200 |
|---|---|---|
| committer | Maze Winther <mazewinther@gmail.com> | 2026-04-07 01:09:13 +0200 |
| commit | e4b67094e757917588fb624a30aa046667403add (patch) | |
| tree | 55452692eac69a5a63b58f2a5cdb701dd7023d68 /rust/crates/masks | |
| parent | 79df73643185c02493364a077402726d1fa7b9ec (diff) | |
| download | opencut-e4b67094e757917588fb624a30aa046667403add.tar.gz | |
refactor: split effects and masks into dedicated rust crates, introduce MediaTime and FrameRate
Diffstat (limited to 'rust/crates/masks')
| -rw-r--r-- | rust/crates/masks/Cargo.toml | 13 | ||||
| -rw-r--r-- | rust/crates/masks/src/feather.rs | 263 | ||||
| -rw-r--r-- | rust/crates/masks/src/masks.rs | 5 | ||||
| -rw-r--r-- | rust/crates/masks/src/sdf.rs | 308 | ||||
| -rw-r--r-- | rust/crates/masks/src/shaders/jfa_distance.wgsl | 50 | ||||
| -rw-r--r-- | rust/crates/masks/src/shaders/jfa_init.wgsl | 35 | ||||
| -rw-r--r-- | rust/crates/masks/src/shaders/jfa_step.wgsl | 75 |
7 files changed, 749 insertions, 0 deletions
diff --git a/rust/crates/masks/Cargo.toml b/rust/crates/masks/Cargo.toml new file mode 100644 index 00000000..58e10712 --- /dev/null +++ b/rust/crates/masks/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "masks" +version = "0.1.0" +edition = "2024" + +[lib] +path = "src/masks.rs" +crate-type = ["rlib"] + +[dependencies] +bytemuck = { version = "1.25.0", features = ["derive"] } +gpu = { version = "0.1.0", path = "../gpu" } +wgpu = "29.0.1" diff --git a/rust/crates/masks/src/feather.rs b/rust/crates/masks/src/feather.rs new file mode 100644 index 00000000..e58572f4 --- /dev/null +++ b/rust/crates/masks/src/feather.rs @@ -0,0 +1,263 @@ +use bytemuck::{Pod, Zeroable}; +use gpu::{FULLSCREEN_SHADER_SOURCE, GPU_TEXTURE_FORMAT, GpuContext}; +use wgpu::util::DeviceExt; + +use crate::SdfPipeline; + +const JFA_DISTANCE_SHADER_SOURCE: &str = include_str!("shaders/jfa_distance.wgsl"); + +pub struct ApplyMaskFeatherOptions<'a> { + pub mask: &'a wgpu::Texture, + pub width: u32, + pub height: u32, + pub feather: f32, +} + +pub struct MaskFeatherPipeline { + sdf_pipeline: SdfPipeline, + inside_texture_bind_group_layout: wgpu::BindGroupLayout, + outside_texture_bind_group_layout: wgpu::BindGroupLayout, + uniform_bind_group_layout: wgpu::BindGroupLayout, + distance_pipeline: wgpu::RenderPipeline, +} + +#[repr(C)] +#[derive(Clone, Copy, Pod, Zeroable)] +struct DistanceUniformBuffer { + resolution: [f32; 2], + feather_half: f32, + _padding: f32, +} + +impl MaskFeatherPipeline { + pub fn new(context: &GpuContext) -> Self { + let device = context.device(); + let inside_texture_bind_group_layout = + device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("gpu-mask-distance-inside-layout"), + entries: &[ + wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: wgpu::ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Texture { + multisampled: false, + view_dimension: wgpu::TextureViewDimension::D2, + sample_type: wgpu::TextureSampleType::Float { filterable: true }, + }, + count: None, + }, + wgpu::BindGroupLayoutEntry { + binding: 1, + visibility: wgpu::ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering), + count: None, + }, + ], + }); + let outside_texture_bind_group_layout = + device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("gpu-mask-distance-outside-layout"), + entries: &[ + wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: wgpu::ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Texture { + multisampled: false, + view_dimension: wgpu::TextureViewDimension::D2, + sample_type: wgpu::TextureSampleType::Float { filterable: true }, + }, + count: None, + }, + wgpu::BindGroupLayoutEntry { + binding: 1, + visibility: wgpu::ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering), + count: None, + }, + ], + }); + let uniform_bind_group_layout = + device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("gpu-mask-distance-uniform-layout"), + entries: &[wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: wgpu::ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Uniform, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }], + }); + let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some("gpu-mask-distance-pipeline-layout"), + bind_group_layouts: &[ + Some(&inside_texture_bind_group_layout), + Some(&outside_texture_bind_group_layout), + Some(&uniform_bind_group_layout), + ], + immediate_size: 0, + }); + let vertex_shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("gpu-mask-distance-fullscreen-shader"), + source: wgpu::ShaderSource::Wgsl(FULLSCREEN_SHADER_SOURCE.into()), + }); + let fragment_shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("gpu-mask-distance-fragment-shader"), + source: wgpu::ShaderSource::Wgsl(JFA_DISTANCE_SHADER_SOURCE.into()), + }); + let distance_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor { + label: Some("gpu-mask-distance-pipeline"), + layout: Some(&pipeline_layout), + vertex: wgpu::VertexState { + module: &vertex_shader_module, + entry_point: Some("vertex_main"), + buffers: &[wgpu::VertexBufferLayout { + array_stride: std::mem::size_of::<[f32; 2]>() as u64, + step_mode: wgpu::VertexStepMode::Vertex, + attributes: &[wgpu::VertexAttribute { + format: wgpu::VertexFormat::Float32x2, + offset: 0, + shader_location: 0, + }], + }], + compilation_options: wgpu::PipelineCompilationOptions::default(), + }, + fragment: Some(wgpu::FragmentState { + module: &fragment_shader_module, + entry_point: Some("fragment_main"), + targets: &[Some(wgpu::ColorTargetState { + format: GPU_TEXTURE_FORMAT, + blend: None, + write_mask: wgpu::ColorWrites::ALL, + })], + compilation_options: wgpu::PipelineCompilationOptions::default(), + }), + primitive: wgpu::PrimitiveState::default(), + depth_stencil: None, + multisample: wgpu::MultisampleState::default(), + multiview_mask: None, + cache: None, + }); + + Self { + sdf_pipeline: SdfPipeline::new(context), + inside_texture_bind_group_layout, + outside_texture_bind_group_layout, + uniform_bind_group_layout, + distance_pipeline, + } + } + + pub fn apply_mask_feather( + &self, + context: &GpuContext, + ApplyMaskFeatherOptions { + mask, + width, + height, + feather, + }: ApplyMaskFeatherOptions<'_>, + ) -> wgpu::Texture { + let sdf = self + .sdf_pipeline + .compute_signed_distance_field(context, mask, width, height); + let output_texture = context.create_render_texture(width, height, "masks-feather-output"); + let inside_view = sdf + .inside_texture + .create_view(&wgpu::TextureViewDescriptor::default()); + let outside_view = sdf + .outside_texture + .create_view(&wgpu::TextureViewDescriptor::default()); + let output_view = output_texture.create_view(&wgpu::TextureViewDescriptor::default()); + let inside_bind_group = context + .device() + .create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("gpu-mask-distance-inside-bind-group"), + layout: &self.inside_texture_bind_group_layout, + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: wgpu::BindingResource::TextureView(&inside_view), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: wgpu::BindingResource::Sampler(context.nearest_sampler()), + }, + ], + }); + let outside_bind_group = context + .device() + .create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("gpu-mask-distance-outside-bind-group"), + layout: &self.outside_texture_bind_group_layout, + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: wgpu::BindingResource::TextureView(&outside_view), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: wgpu::BindingResource::Sampler(context.nearest_sampler()), + }, + ], + }); + let uniform_buffer = + context + .device() + .create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("gpu-mask-distance-uniform-buffer"), + contents: bytemuck::bytes_of(&DistanceUniformBuffer { + resolution: [width as f32, height as f32], + feather_half: feather / 2.0, + _padding: 0.0, + }), + usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST, + }); + let uniform_bind_group = context + .device() + .create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("gpu-mask-distance-uniform-bind-group"), + layout: &self.uniform_bind_group_layout, + entries: &[wgpu::BindGroupEntry { + binding: 0, + resource: uniform_buffer.as_entire_binding(), + }], + }); + let mut encoder = + context + .device() + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("gpu-mask-distance-command-encoder"), + }); + + { + let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor { + label: Some("gpu-mask-distance-render-pass"), + color_attachments: &[Some(wgpu::RenderPassColorAttachment { + view: &output_view, + resolve_target: None, + depth_slice: None, + ops: wgpu::Operations { + load: wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT), + store: wgpu::StoreOp::Store, + }, + })], + depth_stencil_attachment: None, + occlusion_query_set: None, + timestamp_writes: None, + multiview_mask: None, + }); + render_pass.set_pipeline(&self.distance_pipeline); + render_pass.set_vertex_buffer(0, context.fullscreen_quad().slice(..)); + render_pass.set_bind_group(0, &inside_bind_group, &[]); + render_pass.set_bind_group(1, &outside_bind_group, &[]); + render_pass.set_bind_group(2, &uniform_bind_group, &[]); + render_pass.draw(0..6, 0..1); + } + + context.queue().submit([encoder.finish()]); + output_texture + } +} diff --git a/rust/crates/masks/src/masks.rs b/rust/crates/masks/src/masks.rs new file mode 100644 index 00000000..2bf33342 --- /dev/null +++ b/rust/crates/masks/src/masks.rs @@ -0,0 +1,5 @@ +mod feather; +mod sdf; + +pub use feather::{ApplyMaskFeatherOptions, MaskFeatherPipeline}; +pub use sdf::{SdfPipeline, SignedDistanceFieldTextures}; diff --git a/rust/crates/masks/src/sdf.rs b/rust/crates/masks/src/sdf.rs new file mode 100644 index 00000000..0e7ec01a --- /dev/null +++ b/rust/crates/masks/src/sdf.rs @@ -0,0 +1,308 @@ +use bytemuck::{Pod, Zeroable}; +use gpu::{FULLSCREEN_SHADER_SOURCE, GPU_TEXTURE_FORMAT, GpuContext}; +use wgpu::util::DeviceExt; + +const JFA_INIT_SHADER_SOURCE: &str = include_str!("shaders/jfa_init.wgsl"); +const JFA_STEP_SHADER_SOURCE: &str = include_str!("shaders/jfa_step.wgsl"); + +pub struct SignedDistanceFieldTextures { + pub inside_texture: wgpu::Texture, + pub outside_texture: wgpu::Texture, +} + +pub struct SdfPipeline { + texture_bind_group_layout: wgpu::BindGroupLayout, + uniform_bind_group_layout: wgpu::BindGroupLayout, + init_pipeline: wgpu::RenderPipeline, + step_pipeline: wgpu::RenderPipeline, +} + +#[repr(C)] +#[derive(Clone, Copy, Pod, Zeroable)] +struct JfaInitUniformBuffer { + resolution: [f32; 2], + invert: f32, + _padding: f32, +} + +#[repr(C)] +#[derive(Clone, Copy, Pod, Zeroable)] +struct JfaStepUniformBuffer { + resolution: [f32; 2], + step_size: f32, + _padding: f32, +} + +impl SdfPipeline { + pub fn new(context: &GpuContext) -> Self { + let device = context.device(); + let texture_bind_group_layout = + device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("gpu-sdf-texture-bind-group-layout"), + entries: &[ + wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: wgpu::ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Texture { + multisampled: false, + view_dimension: wgpu::TextureViewDimension::D2, + sample_type: wgpu::TextureSampleType::Float { filterable: true }, + }, + count: None, + }, + wgpu::BindGroupLayoutEntry { + binding: 1, + visibility: wgpu::ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering), + count: None, + }, + ], + }); + let uniform_bind_group_layout = + device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("gpu-sdf-uniform-bind-group-layout"), + entries: &[wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: wgpu::ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Uniform, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }], + }); + let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some("gpu-sdf-pipeline-layout"), + bind_group_layouts: &[ + Some(&texture_bind_group_layout), + Some(&uniform_bind_group_layout), + ], + immediate_size: 0, + }); + let vertex_shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("gpu-sdf-fullscreen-shader"), + source: wgpu::ShaderSource::Wgsl(FULLSCREEN_SHADER_SOURCE.into()), + }); + let init_shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("gpu-jfa-init-shader"), + source: wgpu::ShaderSource::Wgsl(JFA_INIT_SHADER_SOURCE.into()), + }); + let step_shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("gpu-jfa-step-shader"), + source: wgpu::ShaderSource::Wgsl(JFA_STEP_SHADER_SOURCE.into()), + }); + let init_pipeline = create_pipeline( + device, + "gpu-jfa-init-pipeline", + &pipeline_layout, + &vertex_shader_module, + &init_shader_module, + ); + let step_pipeline = create_pipeline( + device, + "gpu-jfa-step-pipeline", + &pipeline_layout, + &vertex_shader_module, + &step_shader_module, + ); + + Self { + texture_bind_group_layout, + uniform_bind_group_layout, + init_pipeline, + step_pipeline, + } + } + + pub fn compute_signed_distance_field( + &self, + context: &GpuContext, + source_texture: &wgpu::Texture, + width: u32, + height: u32, + ) -> SignedDistanceFieldTextures { + SignedDistanceFieldTextures { + inside_texture: self.run_jfa(context, source_texture, width, height, false), + outside_texture: self.run_jfa(context, source_texture, width, height, true), + } + } + + fn run_jfa( + &self, + context: &GpuContext, + source_texture: &wgpu::Texture, + width: u32, + height: u32, + is_inverted: bool, + ) -> wgpu::Texture { + let ping_texture = context.create_render_texture(width, height, "gpu-jfa-ping-texture"); + let pong_texture = context.create_render_texture(width, height, "gpu-jfa-pong-texture"); + + self.run_pass( + context, + source_texture, + &ping_texture, + &self.init_pipeline, + bytemuck::bytes_of(&JfaInitUniformBuffer { + resolution: [width as f32, height as f32], + invert: if is_inverted { 1.0 } else { 0.0 }, + _padding: 0.0, + }), + ); + + let mut source_is_ping = true; + let steps = (width.max(height) as f32).log2().ceil() as u32; + for step_index in (0..steps).rev() { + let step_size = 2u32.pow(step_index).max(1); + let input_texture = if source_is_ping { + &ping_texture + } else { + &pong_texture + }; + let output_texture = if source_is_ping { + &pong_texture + } else { + &ping_texture + }; + self.run_pass( + context, + input_texture, + output_texture, + &self.step_pipeline, + bytemuck::bytes_of(&JfaStepUniformBuffer { + resolution: [width as f32, height as f32], + step_size: step_size as f32, + _padding: 0.0, + }), + ); + source_is_ping = !source_is_ping; + } + + if source_is_ping { + ping_texture + } else { + pong_texture + } + } + + fn run_pass( + &self, + context: &GpuContext, + input_texture: &wgpu::Texture, + output_texture: &wgpu::Texture, + pipeline: &wgpu::RenderPipeline, + uniform_buffer_bytes: &[u8], + ) { + let input_view = input_texture.create_view(&wgpu::TextureViewDescriptor::default()); + let output_view = output_texture.create_view(&wgpu::TextureViewDescriptor::default()); + let texture_bind_group = context + .device() + .create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("gpu-sdf-texture-bind-group"), + layout: &self.texture_bind_group_layout, + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: wgpu::BindingResource::TextureView(&input_view), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: wgpu::BindingResource::Sampler(context.nearest_sampler()), + }, + ], + }); + let uniform_buffer = + context + .device() + .create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("gpu-sdf-uniform-buffer"), + contents: uniform_buffer_bytes, + usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST, + }); + let uniform_bind_group = context + .device() + .create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("gpu-sdf-uniform-bind-group"), + layout: &self.uniform_bind_group_layout, + entries: &[wgpu::BindGroupEntry { + binding: 0, + resource: uniform_buffer.as_entire_binding(), + }], + }); + let mut encoder = + context + .device() + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("gpu-sdf-command-encoder"), + }); + + { + let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor { + label: Some("gpu-sdf-render-pass"), + color_attachments: &[Some(wgpu::RenderPassColorAttachment { + view: &output_view, + resolve_target: None, + depth_slice: None, + ops: wgpu::Operations { + load: wgpu::LoadOp::Clear(wgpu::Color::WHITE), + store: wgpu::StoreOp::Store, + }, + })], + depth_stencil_attachment: None, + occlusion_query_set: None, + timestamp_writes: None, + multiview_mask: None, + }); + render_pass.set_pipeline(pipeline); + render_pass.set_vertex_buffer(0, context.fullscreen_quad().slice(..)); + render_pass.set_bind_group(0, &texture_bind_group, &[]); + render_pass.set_bind_group(1, &uniform_bind_group, &[]); + render_pass.draw(0..6, 0..1); + } + + context.queue().submit([encoder.finish()]); + } +} + +fn create_pipeline( + device: &wgpu::Device, + label: &'static str, + layout: &wgpu::PipelineLayout, + vertex_shader_module: &wgpu::ShaderModule, + fragment_shader_module: &wgpu::ShaderModule, +) -> wgpu::RenderPipeline { + device.create_render_pipeline(&wgpu::RenderPipelineDescriptor { + label: Some(label), + layout: Some(layout), + vertex: wgpu::VertexState { + module: vertex_shader_module, + entry_point: Some("vertex_main"), + buffers: &[wgpu::VertexBufferLayout { + array_stride: std::mem::size_of::<[f32; 2]>() as u64, + step_mode: wgpu::VertexStepMode::Vertex, + attributes: &[wgpu::VertexAttribute { + format: wgpu::VertexFormat::Float32x2, + offset: 0, + shader_location: 0, + }], + }], + compilation_options: wgpu::PipelineCompilationOptions::default(), + }, + fragment: Some(wgpu::FragmentState { + module: fragment_shader_module, + entry_point: Some("fragment_main"), + targets: &[Some(wgpu::ColorTargetState { + format: GPU_TEXTURE_FORMAT, + blend: None, + write_mask: wgpu::ColorWrites::ALL, + })], + compilation_options: wgpu::PipelineCompilationOptions::default(), + }), + primitive: wgpu::PrimitiveState::default(), + depth_stencil: None, + multisample: wgpu::MultisampleState::default(), + multiview_mask: None, + cache: None, + }) +} diff --git a/rust/crates/masks/src/shaders/jfa_distance.wgsl b/rust/crates/masks/src/shaders/jfa_distance.wgsl new file mode 100644 index 00000000..e65a7023 --- /dev/null +++ b/rust/crates/masks/src/shaders/jfa_distance.wgsl @@ -0,0 +1,50 @@ +struct VertexOutput { + @builtin(position) position: vec4f, + @location(0) tex_coord: vec2f, +} + +struct DistanceUniforms { + resolution: vec2f, + feather_half: f32, + _padding: f32, +} + +@group(0) @binding(0) var inside_texture: texture_2d<f32>; +@group(0) @binding(1) var inside_sampler: sampler; +@group(1) @binding(0) var outside_texture: texture_2d<f32>; +@group(1) @binding(1) var outside_sampler: sampler; +@group(2) @binding(0) var<uniform> uniforms: DistanceUniforms; + +fn decode_seed(encoded: vec4f) -> vec2f { + let x = floor(encoded.r * 255.0 + 0.5) * 256.0 + floor(encoded.g * 255.0 + 0.5); + let y = floor(encoded.b * 255.0 + 0.5) * 256.0 + floor(encoded.a * 255.0 + 0.5); + return vec2f(x, y); +} + +fn is_no_seed(encoded: vec4f) -> bool { + return encoded.r > 0.99 && encoded.g > 0.99 && encoded.b > 0.99 && encoded.a > 0.99; +} + +@fragment +fn fragment_main(input: VertexOutput) -> @location(0) vec4f { + let pixel_coord = floor(input.tex_coord * uniforms.resolution); + let inside_encoded = textureSample(inside_texture, inside_sampler, input.tex_coord); + let outside_encoded = textureSample(outside_texture, outside_sampler, input.tex_coord); + + let has_inside = !is_no_seed(inside_encoded); + let has_outside = !is_no_seed(outside_encoded); + let distance_to_inside = select( + 100000.0, + distance(pixel_coord, decode_seed(inside_encoded)), + has_inside, + ); + let distance_to_outside = select( + 100000.0, + distance(pixel_coord, decode_seed(outside_encoded)), + has_outside, + ); + let signed_distance = distance_to_outside - distance_to_inside; + let alpha = smoothstep(-uniforms.feather_half, uniforms.feather_half, signed_distance); + + return vec4f(alpha, alpha, alpha, alpha); +} diff --git a/rust/crates/masks/src/shaders/jfa_init.wgsl b/rust/crates/masks/src/shaders/jfa_init.wgsl new file mode 100644 index 00000000..29aa38c7 --- /dev/null +++ b/rust/crates/masks/src/shaders/jfa_init.wgsl @@ -0,0 +1,35 @@ +struct VertexOutput { + @builtin(position) position: vec4f, + @location(0) tex_coord: vec2f, +} + +struct JfaInitUniforms { + resolution: vec2f, + invert: f32, + _padding: f32, +} + +@group(0) @binding(0) var input_texture: texture_2d<f32>; +@group(0) @binding(1) var input_sampler: sampler; +@group(1) @binding(0) var<uniform> uniforms: JfaInitUniforms; + +fn encode_seed(seed: vec2f) -> vec4f { + let x_hi = floor(seed.x / 256.0); + let x_lo = seed.x - (x_hi * 256.0); + let y_hi = floor(seed.y / 256.0); + let y_lo = seed.y - (y_hi * 256.0); + return vec4f(x_hi / 255.0, x_lo / 255.0, y_hi / 255.0, y_lo / 255.0); +} + +@fragment +fn fragment_main(input: VertexOutput) -> @location(0) vec4f { + let mask = textureSample(input_texture, input_sampler, input.tex_coord).r; + let is_seed = select(mask > 0.5, mask < 0.5, uniforms.invert > 0.5); + + if (is_seed) { + let pixel_coord = floor(input.tex_coord * uniforms.resolution); + return encode_seed(pixel_coord); + } + + return vec4f(1.0, 1.0, 1.0, 1.0); +} diff --git a/rust/crates/masks/src/shaders/jfa_step.wgsl b/rust/crates/masks/src/shaders/jfa_step.wgsl new file mode 100644 index 00000000..e20d7fb0 --- /dev/null +++ b/rust/crates/masks/src/shaders/jfa_step.wgsl @@ -0,0 +1,75 @@ +struct VertexOutput { + @builtin(position) position: vec4f, + @location(0) tex_coord: vec2f, +} + +struct JfaStepUniforms { + resolution: vec2f, + step_size: f32, + _padding: f32, +} + +@group(0) @binding(0) var input_texture: texture_2d<f32>; +@group(0) @binding(1) var input_sampler: sampler; +@group(1) @binding(0) var<uniform> uniforms: JfaStepUniforms; + +fn decode_seed(encoded: vec4f) -> vec2f { + let x = floor(encoded.r * 255.0 + 0.5) * 256.0 + floor(encoded.g * 255.0 + 0.5); + let y = floor(encoded.b * 255.0 + 0.5) * 256.0 + floor(encoded.a * 255.0 + 0.5); + return vec2f(x, y); +} + +fn encode_seed(seed: vec2f) -> vec4f { + let x_hi = floor(seed.x / 256.0); + let x_lo = seed.x - (x_hi * 256.0); + let y_hi = floor(seed.y / 256.0); + let y_lo = seed.y - (y_hi * 256.0); + return vec4f(x_hi / 255.0, x_lo / 255.0, y_hi / 255.0, y_lo / 255.0); +} + +fn is_no_seed(encoded: vec4f) -> bool { + return encoded.r > 0.99 && encoded.g > 0.99 && encoded.b > 0.99 && encoded.a > 0.99; +} + +@fragment +fn fragment_main(input: VertexOutput) -> @location(0) vec4f { + let pixel_coord = floor(input.tex_coord * uniforms.resolution); + let texel_size = vec2f(1.0, 1.0) / uniforms.resolution; + + var best_distance = 10000000000.0; + var best_seed = vec2f(65535.0, 65535.0); + + for (var y = -1; y <= 1; y = y + 1) { + for (var x = -1; x <= 1; x = x + 1) { + let offset = vec2f(f32(x), f32(y)) * uniforms.step_size; + let sample_uv = input.tex_coord + (offset * texel_size); + + if ( + sample_uv.x < 0.0 || + sample_uv.x > 1.0 || + sample_uv.y < 0.0 || + sample_uv.y > 1.0 + ) { + continue; + } + + let encoded = textureSample(input_texture, input_sampler, sample_uv); + if (is_no_seed(encoded)) { + continue; + } + + let seed = decode_seed(encoded); + let distance_to_seed = distance(pixel_coord, seed); + if (distance_to_seed < best_distance) { + best_distance = distance_to_seed; + best_seed = seed; + } + } + } + + if (best_distance < 1000000000.0) { + return encode_seed(best_seed); + } + + return vec4f(1.0, 1.0, 1.0, 1.0); +} |
