1use wgpu::util::DeviceExt;
2
3pub const MAX_RADIUS: f32 = 32.0;
4pub const TILE: u32 = 16;
5const REACH: i32 = 2;
6const MAX_RINGS: u32 = 5;
7const RING_SPACING: f32 = 6.0;
8const RING_TURN: f32 = 2.399_963;
9
10#[derive(Clone, Copy, Debug, PartialEq)]
11pub struct Lens {
12 pub focal_m: f32,
13 pub sensor_m: f32,
14 pub fstop: f32,
15 pub aperture: f32,
16 pub focus_m: f32,
17 pub near_m: f32,
18 pub near_blur: f32,
19 pub far_m: f32,
20 pub far_blur: f32,
21 pub ground_m: Option<f32>,
22}
23
24impl Lens {
25 pub fn enabled(self) -> bool {
26 self.aperture > 0.0 && self.focus_m > self.focal_m && self.sensor_m > 0.0
27 }
28
29 pub fn sensor_for_view(focal_m: f32, focus_m: f32, tan_half_width: f32) -> f32 {
30 2.0 * tan_half_width * focal_m * focus_m / (focus_m - focal_m).max(1.0e-6)
31 }
32
33 pub fn thin_lens(self, distance: f32, width: u32) -> f32 {
34 if !self.enabled() || distance <= 0.0 {
35 return 0.0;
36 }
37 let diameter = self.focal_m * self.aperture / self.fstop.max(0.01);
38 let blur = diameter * self.focal_m / (self.focus_m - self.focal_m)
39 * (distance - self.focus_m)
40 / distance;
41 blur * width as f32 / (2.0 * self.sensor_m.max(1.0e-4))
42 }
43
44 pub fn spreads(self, height: f32) -> bool {
45 self.ground_m.is_none_or(|ground| height < ground)
46 }
47
48 pub fn opening(self, distance: f32, height: f32) -> f32 {
49 let near = if self.spreads(height) {
50 self.near_blur.max(0.0) * (self.near_m - distance).max(0.0) / self.near_m.max(0.001)
51 } else {
52 0.0
53 };
54 let far = self.far_blur.max(0.0) * (distance - self.far_m).max(0.0) / self.far_m.max(0.001);
55 1.0 + near + far
56 }
57
58 pub fn circle_of_confusion(self, distance: f32, height: f32, width: u32) -> f32 {
59 (self.thin_lens(distance, width) * self.opening(distance, height))
60 .clamp(-MAX_RADIUS, MAX_RADIUS)
61 }
62}
63
64pub fn height_plane(view: [[f32; 4]; 4], projection: [[f32; 4]; 4]) -> [f32; 4] {
65 let row = [view[1][0], view[1][1], view[1][2]];
66 let x = row[0] / projection[0][0];
67 let y = row[1] / projection[1][1];
68 let z = row[0] * projection[2][0] / projection[0][0]
69 + row[1] * projection[2][1] / projection[1][1]
70 - row[2];
71 let w = -(row[0] * view[3][0] + row[1] * view[3][1] + row[2] * view[3][2]);
72 [x, y, z, w]
73}
74
75fn slack(distance: f32) -> f32 {
76 0.005 + 0.02 * distance
77}
78
79fn rings(radius: f32) -> u32 {
80 ((radius / RING_SPACING).ceil() as u32).clamp(1, MAX_RINGS)
81}
82
83pub fn gather(
84 image: &[[f32; 4]],
85 depth: &[f32],
86 heights: &[f32],
87 width: usize,
88 height: usize,
89 lens: Lens,
90) -> Vec<[f32; 4]> {
91 let count = width * height;
92 if !lens.enabled() || image.len() != count || depth.len() != count || heights.len() != count {
93 return image.to_vec();
94 }
95 let coc: Vec<f32> = (0..count)
96 .map(|i| lens.circle_of_confusion(depth[i], heights[i], width as u32))
97 .collect();
98 let tiles_x = width.div_ceil(TILE as usize);
99 let tiles_y = height.div_ceil(TILE as usize);
100 let mut tiles = vec![(0.0f32, f32::INFINITY); tiles_x * tiles_y];
101 for y in 0..height {
102 for x in 0..width {
103 let tile = &mut tiles[(y / TILE as usize) * tiles_x + x / TILE as usize];
104 tile.0 = tile.0.max(coc[y * width + x].abs());
105 tile.1 = tile.1.min(depth[y * width + x]);
106 }
107 }
108 let dilated: Vec<(f32, f32)> = (0..tiles_x * tiles_y)
109 .map(|index| {
110 let (tx, ty) = ((index % tiles_x) as i32, (index / tiles_x) as i32);
111 let mut out = (0.0f32, f32::INFINITY);
112 for dy in -REACH..=REACH {
113 for dx in -REACH..=REACH {
114 let (x, y) = (tx + dx, ty + dy);
115 if x >= 0 && y >= 0 && (x as usize) < tiles_x && (y as usize) < tiles_y {
116 let tile = tiles[y as usize * tiles_x + x as usize];
117 out = (out.0.max(tile.0), out.1.min(tile.1));
118 }
119 }
120 }
121 out
122 })
123 .collect();
124 let mut out = image.to_vec();
125 for y in 0..height {
126 for x in 0..width {
127 let index = y * width + x;
128 let d0 = depth[index];
129 let c0 = coc[index];
130 let surface = lens.spreads(heights[index]);
131 let tile = dilated[(y / TILE as usize) * tiles_x + x / TILE as usize];
132 let eps = slack(d0);
133 let front = if tile.1 < d0 - eps { tile.0 } else { 0.0 };
134 let radius = front
135 .max(if surface { c0.abs() } else { 0.0 })
136 .min(MAX_RADIUS);
137 if radius < 1.0 {
138 continue;
139 }
140 let ring_count = rings(radius);
141 let share = radius * radius / (4 * ring_count * (ring_count + 1)) as f32;
142 let centre = [x as f32 + 0.5, y as f32 + 0.5];
143 let mut back = image[index];
144 let mut back_weight = 1.0;
145 let mut fore = [0.0f32; 4];
146 let mut fore_weight = 0.0;
147 for ring in 1..=ring_count {
148 let reach = radius * ring as f32 / ring_count as f32;
149 let taps = 8 * ring;
150 for tap in 0..taps {
151 let angle =
152 RING_TURN * ring as f32 + std::f32::consts::TAU * tap as f32 / taps as f32;
153 let at = [
154 centre[0] + angle.cos() * reach,
155 centre[1] + angle.sin() * reach,
156 ];
157 let sx = (at[0].floor() as i64).clamp(0, width as i64 - 1) as usize;
158 let sy = (at[1].floor() as i64).clamp(0, height as i64 - 1) as usize;
159 let sample = sy * width + sx;
160 let ds = depth[sample];
161 let cs = coc[sample];
162 let colour = image[sample];
163 if ds < d0 - eps {
164 let cover = (cs.abs() - reach + 0.5).clamp(0.0, 1.0);
165 let weight = cover * share / (cs * cs).max(1.0);
166 for c in 0..4 {
167 fore[c] += colour[c] * weight;
168 }
169 fore_weight += weight;
170 } else if surface {
171 let limit = if ds <= d0 + eps {
172 c0.abs().min(cs.abs())
173 } else {
174 c0.abs()
175 };
176 let cover = (limit - reach + 0.5).clamp(0.0, 1.0);
177 for c in 0..4 {
178 back[c] += colour[c] * cover;
179 }
180 back_weight += cover;
181 }
182 }
183 }
184 let alpha = fore_weight.min(1.0);
185 for c in 0..4 {
186 let behind = back[c] / back_weight;
187 let ahead = fore[c] / fore_weight.max(1.0e-5);
188 out[index][c] = behind + (ahead - behind) * alpha;
189 }
190 }
191 }
192 out
193}
194
195const COMMON: &str = r#"
196struct Params {
197 size: vec2u, tiles: vec2u,
198 optics: vec4f, focus: vec4f, far: vec4f, height: vec4f,
199}
200@group(0) @binding(0) var<uniform> lens: Params;
201fn world_height(pixel: vec2f, distance: f32) -> f32 {
202 let ndc = vec2f(pixel.x / f32(lens.size.x) * 2.0 - 1.0, 1.0 - pixel.y / f32(lens.size.y) * 2.0);
203 return dot(lens.height.xyz, vec3f(ndc.x * distance, ndc.y * distance, distance)) + lens.height.w;
204}
205fn spreads(height: f32) -> bool {
206 return lens.far.z < 0.5 || height < lens.far.y;
207}
208fn coc(distance: f32, height: f32) -> f32 {
209 let focal = lens.optics.x;
210 let focus = lens.focus.x;
211 if (distance <= 0.0 || focus <= focal) { return 0.0; }
212 let diameter = focal * lens.optics.w / max(lens.optics.z, 0.01);
213 let thin = diameter * focal / (focus - focal) * (distance - focus) / distance;
214 let low = select(0.0, 1.0, spreads(height));
215 let open = 1.0 + low * lens.focus.z * max(lens.focus.y - distance, 0.0) / max(lens.focus.y, 0.001)
216 + lens.far.x * max(distance - lens.focus.w, 0.0) / max(lens.focus.w, 0.001);
217 return clamp(thin * open * f32(lens.size.x) / (2.0 * max(lens.optics.y, 0.0001)), -lens.far.w, lens.far.w);
218}
219fn slack(distance: f32) -> f32 {
220 return 0.005 + 0.02 * distance;
221}
222"#;
223
224const TILE_WGSL: &str = r#"
225@group(0) @binding(1) var depth: texture_2d<f32>;
226@group(0) @binding(2) var tiles: texture_storage_2d<rgba16float, write>;
227var<workgroup> spread: atomic<u32>;
228var<workgroup> nearest: atomic<u32>;
229@compute @workgroup_size(16, 16)
230fn main(@builtin(global_invocation_id) id: vec3u, @builtin(local_invocation_index) local: u32, @builtin(workgroup_id) group: vec3u) {
231 if (local == 0u) {
232 atomicStore(&spread, 0u);
233 atomicStore(&nearest, bitcast<u32>(60000.0));
234 }
235 workgroupBarrier();
236 if (all(id.xy < lens.size)) {
237 let distance = textureLoad(depth, vec2i(id.xy), 0).x;
238 let height = world_height(vec2f(id.xy) + vec2f(0.5), distance);
239 atomicMax(&spread, bitcast<u32>(abs(coc(distance, height))));
240 atomicMin(&nearest, bitcast<u32>(clamp(distance, 0.0, 60000.0)));
241 }
242 workgroupBarrier();
243 if (local == 0u) {
244 textureStore(tiles, vec2i(group.xy), vec4f(bitcast<f32>(atomicLoad(&spread)), bitcast<f32>(atomicLoad(&nearest)), 0.0, 0.0));
245 }
246}
247"#;
248
249const DILATE_WGSL: &str = r#"
250@group(0) @binding(1) var tiles: texture_2d<f32>;
251@group(0) @binding(2) var dilated: texture_storage_2d<rgba16float, write>;
252@compute @workgroup_size(8, 8)
253fn main(@builtin(global_invocation_id) id: vec3u) {
254 if (any(id.xy >= lens.tiles)) { return; }
255 var spread = 0.0;
256 var nearest = 60000.0;
257 for (var dy = -2; dy <= 2; dy++) {
258 for (var dx = -2; dx <= 2; dx++) {
259 let at = vec2i(id.xy) + vec2i(dx, dy);
260 if (any(at < vec2i(0)) || any(at >= vec2i(lens.tiles))) { continue; }
261 let tile = textureLoad(tiles, at, 0);
262 spread = max(spread, tile.x);
263 nearest = min(nearest, tile.y);
264 }
265 }
266 textureStore(dilated, vec2i(id.xy), vec4f(spread, nearest, 0.0, 0.0));
267}
268"#;
269
270const GATHER_WGSL: &str = r#"
271@group(0) @binding(1) var source: texture_2d<f32>;
272@group(0) @binding(2) var depth: texture_2d<f32>;
273@group(0) @binding(3) var dilated: texture_2d<f32>;
274@group(0) @binding(4) var output_tex: texture_storage_2d<rgba16float, write>;
275@compute @workgroup_size(8, 8)
276fn main(@builtin(global_invocation_id) id: vec3u) {
277 if (any(id.xy >= lens.size)) { return; }
278 let pixel = vec2i(id.xy);
279 let sharp = textureLoad(source, pixel, 0);
280 let d0 = textureLoad(depth, pixel, 0).x;
281 let centre = vec2f(pixel) + vec2f(0.5);
282 let c0 = coc(d0, world_height(centre, d0));
283 let surface = spreads(world_height(centre, d0));
284 let tile = textureLoad(dilated, pixel / 16, 0);
285 let eps = slack(d0);
286 let front = select(0.0, tile.x, tile.y < d0 - eps);
287 let radius = min(max(front, select(0.0, abs(c0), surface)), lens.far.w);
288 if (radius < 1.0) {
289 textureStore(output_tex, pixel, sharp);
290 return;
291 }
292 let rings = clamp(u32(ceil(radius / 6.0)), 1u, 5u);
293 let share = radius * radius / f32(4u * rings * (rings + 1u));
294 let last = vec2i(lens.size) - vec2i(1);
295 var back = sharp;
296 var back_weight = 1.0;
297 var fore = vec4f(0.0);
298 var fore_weight = 0.0;
299 for (var ring = 1u; ring <= rings; ring++) {
300 let reach = radius * f32(ring) / f32(rings);
301 let taps = 8u * ring;
302 for (var tap = 0u; tap < taps; tap++) {
303 let angle = 2.399963 * f32(ring) + 6.2831853 * f32(tap) / f32(taps);
304 let at = centre + vec2f(cos(angle), sin(angle)) * reach;
305 let texel = clamp(vec2i(floor(at)), vec2i(0), last);
306 let ds = textureLoad(depth, texel, 0).x;
307 let cs = coc(ds, world_height(vec2f(texel) + vec2f(0.5), ds));
308 let colour = textureLoad(source, texel, 0);
309 if (ds < d0 - eps) {
310 let cover = clamp(abs(cs) - reach + 0.5, 0.0, 1.0);
311 let weight = cover * share / max(cs * cs, 1.0);
312 fore += colour * weight;
313 fore_weight += weight;
314 } else if (surface) {
315 let limit = select(abs(c0), min(abs(c0), abs(cs)), ds <= d0 + eps);
316 let cover = clamp(limit - reach + 0.5, 0.0, 1.0);
317 back += colour * cover;
318 back_weight += cover;
319 }
320 }
321 }
322 let behind = back / back_weight;
323 let ahead = fore / max(fore_weight, 0.00001);
324 textureStore(output_tex, pixel, mix(behind, ahead, min(fore_weight, 1.0)));
325}
326"#;
327
328fn shader(part: &str) -> String {
329 format!("{COMMON}{part}")
330}
331
332fn uniform_entry() -> wgpu::BindGroupLayoutEntry {
333 wgpu::BindGroupLayoutEntry {
334 binding: 0,
335 visibility: wgpu::ShaderStages::COMPUTE,
336 ty: wgpu::BindingType::Buffer {
337 ty: wgpu::BufferBindingType::Uniform,
338 has_dynamic_offset: false,
339 min_binding_size: None,
340 },
341 count: None,
342 }
343}
344
345fn texture_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
346 wgpu::BindGroupLayoutEntry {
347 binding,
348 visibility: wgpu::ShaderStages::COMPUTE,
349 ty: wgpu::BindingType::Texture {
350 sample_type: wgpu::TextureSampleType::Float { filterable: false },
351 view_dimension: wgpu::TextureViewDimension::D2,
352 multisampled: false,
353 },
354 count: None,
355 }
356}
357
358fn storage_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
359 wgpu::BindGroupLayoutEntry {
360 binding,
361 visibility: wgpu::ShaderStages::COMPUTE,
362 ty: wgpu::BindingType::StorageTexture {
363 access: wgpu::StorageTextureAccess::WriteOnly,
364 format: wgpu::TextureFormat::Rgba16Float,
365 view_dimension: wgpu::TextureViewDimension::D2,
366 },
367 count: None,
368 }
369}
370
371fn compute(
372 device: &wgpu::Device,
373 label: &'static str,
374 source: &str,
375 entries: &[wgpu::BindGroupLayoutEntry],
376) -> (wgpu::BindGroupLayout, wgpu::ComputePipeline) {
377 let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
378 label: Some(label),
379 entries,
380 });
381 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
382 label: Some(label),
383 bind_group_layouts: &[Some(&layout)],
384 immediate_size: 0,
385 });
386 let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
387 label: Some(label),
388 source: wgpu::ShaderSource::Wgsl(shader(source).into()),
389 });
390 let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
391 label: Some(label),
392 layout: Some(&pipeline_layout),
393 module: &module,
394 entry_point: Some("main"),
395 compilation_options: Default::default(),
396 cache: None,
397 });
398 (layout, pipeline)
399}
400
401struct Targets {
402 _output: wgpu::Texture,
403 output: wgpu::TextureView,
404 _tiles: wgpu::Texture,
405 tiles: wgpu::TextureView,
406 _dilated: wgpu::Texture,
407 dilated: wgpu::TextureView,
408}
409
410impl Targets {
411 fn new(device: &wgpu::Device, width: u32, height: u32) -> Self {
412 let texture = |label, width, height| {
413 let texture = device.create_texture(&wgpu::TextureDescriptor {
414 label: Some(label),
415 size: wgpu::Extent3d {
416 width,
417 height,
418 depth_or_array_layers: 1,
419 },
420 mip_level_count: 1,
421 sample_count: 1,
422 dimension: wgpu::TextureDimension::D2,
423 format: wgpu::TextureFormat::Rgba16Float,
424 usage: wgpu::TextureUsages::TEXTURE_BINDING
425 | wgpu::TextureUsages::STORAGE_BINDING
426 | wgpu::TextureUsages::COPY_SRC,
427 view_formats: &[],
428 });
429 let view = texture.create_view(&Default::default());
430 (texture, view)
431 };
432 let (tiles_x, tiles_y) = (width.div_ceil(TILE), height.div_ceil(TILE));
433 let (output_texture, output) = texture("depth of field target", width, height);
434 let (tiles_texture, tiles) = texture("depth of field tiles", tiles_x, tiles_y);
435 let (dilated_texture, dilated) = texture("depth of field dilated tiles", tiles_x, tiles_y);
436 Self {
437 _output: output_texture,
438 output,
439 _tiles: tiles_texture,
440 tiles,
441 _dilated: dilated_texture,
442 dilated,
443 }
444 }
445}
446
447pub struct GpuDof {
448 device: wgpu::Device,
449 width: u32,
450 height: u32,
451 viewport: (u32, u32),
452 targets: Targets,
453 tile_layout: wgpu::BindGroupLayout,
454 tile_pipeline: wgpu::ComputePipeline,
455 dilate_layout: wgpu::BindGroupLayout,
456 dilate_pipeline: wgpu::ComputePipeline,
457 gather_layout: wgpu::BindGroupLayout,
458 gather_pipeline: wgpu::ComputePipeline,
459}
460
461impl GpuDof {
462 pub fn new(device: &wgpu::Device, width: u32, height: u32) -> Self {
463 let (tile_layout, tile_pipeline) = compute(
464 device,
465 "depth of field tiles",
466 TILE_WGSL,
467 &[uniform_entry(), texture_entry(1), storage_entry(2)],
468 );
469 let (dilate_layout, dilate_pipeline) = compute(
470 device,
471 "depth of field dilate",
472 DILATE_WGSL,
473 &[uniform_entry(), texture_entry(1), storage_entry(2)],
474 );
475 let (gather_layout, gather_pipeline) = compute(
476 device,
477 "depth of field gather",
478 GATHER_WGSL,
479 &[
480 uniform_entry(),
481 texture_entry(1),
482 texture_entry(2),
483 texture_entry(3),
484 storage_entry(4),
485 ],
486 );
487 Self {
488 device: device.clone(),
489 width,
490 height,
491 viewport: (width, height),
492 targets: Targets::new(device, width, height),
493 tile_layout,
494 tile_pipeline,
495 dilate_layout,
496 dilate_pipeline,
497 gather_layout,
498 gather_pipeline,
499 }
500 }
501
502 pub fn resize(&mut self, width: u32, height: u32) {
503 if (width, height) != (self.width, self.height) {
504 self.targets = Targets::new(&self.device, width, height);
505 self.width = width;
506 self.height = height;
507 }
508 self.viewport = (width, height);
509 }
510
511 pub fn capacity(&self) -> (u32, u32) {
512 (self.width, self.height)
513 }
514
515 pub fn viewport(&self) -> (u32, u32) {
516 self.viewport
517 }
518
519 pub fn set_viewport(&mut self, width: u32, height: u32) -> Result<(), String> {
520 self.viewport = crate::gpu::fit_viewport((width, height), (self.width, self.height))?;
521 Ok(())
522 }
523
524 fn params(&self, lens: Lens, height: [f32; 4]) -> [f32; 20] {
525 let (width, rows) = self.viewport;
526 let (tiles_x, tiles_y) = (width.div_ceil(TILE), rows.div_ceil(TILE));
527 [
528 f32::from_bits(width),
529 f32::from_bits(rows),
530 f32::from_bits(tiles_x),
531 f32::from_bits(tiles_y),
532 lens.focal_m,
533 lens.sensor_m,
534 lens.fstop,
535 lens.aperture,
536 lens.focus_m,
537 lens.near_m,
538 lens.near_blur,
539 lens.far_m,
540 lens.far_blur,
541 lens.ground_m.unwrap_or(0.0),
542 if lens.ground_m.is_some() { 1.0 } else { 0.0 },
543 MAX_RADIUS,
544 height[0],
545 height[1],
546 height[2],
547 height[3],
548 ]
549 }
550
551 fn bind(
552 &self,
553 label: &'static str,
554 layout: &wgpu::BindGroupLayout,
555 resources: &[wgpu::BindingResource<'_>],
556 ) -> wgpu::BindGroup {
557 let entries: Vec<_> = resources
558 .iter()
559 .enumerate()
560 .map(|(binding, resource)| wgpu::BindGroupEntry {
561 binding: binding as u32,
562 resource: resource.clone(),
563 })
564 .collect();
565 self.device.create_bind_group(&wgpu::BindGroupDescriptor {
566 label: Some(label),
567 layout,
568 entries: &entries,
569 })
570 }
571
572 pub fn render<'a>(
573 &'a self,
574 encoder: &mut wgpu::CommandEncoder,
575 source: &wgpu::TextureView,
576 depth: &wgpu::TextureView,
577 lens: Lens,
578 height: [f32; 4],
579 timestamps: Option<wgpu::ComputePassTimestampWrites<'_>>,
580 ) -> &'a wgpu::TextureView {
581 let uniform = self
582 .device
583 .create_buffer_init(&wgpu::util::BufferInitDescriptor {
584 label: Some("depth of field lens"),
585 contents: bytemuck::cast_slice(&self.params(lens, height)),
586 usage: wgpu::BufferUsages::UNIFORM,
587 });
588 let targets = &self.targets;
589 let tiles = self.bind(
590 "depth of field tile bindings",
591 &self.tile_layout,
592 &[
593 uniform.as_entire_binding(),
594 wgpu::BindingResource::TextureView(depth),
595 wgpu::BindingResource::TextureView(&targets.tiles),
596 ],
597 );
598 let dilate = self.bind(
599 "depth of field dilate bindings",
600 &self.dilate_layout,
601 &[
602 uniform.as_entire_binding(),
603 wgpu::BindingResource::TextureView(&targets.tiles),
604 wgpu::BindingResource::TextureView(&targets.dilated),
605 ],
606 );
607 let gather = self.bind(
608 "depth of field gather bindings",
609 &self.gather_layout,
610 &[
611 uniform.as_entire_binding(),
612 wgpu::BindingResource::TextureView(source),
613 wgpu::BindingResource::TextureView(depth),
614 wgpu::BindingResource::TextureView(&targets.dilated),
615 wgpu::BindingResource::TextureView(&targets.output),
616 ],
617 );
618 let (width, rows) = self.viewport;
619 let (tiles_x, tiles_y) = (width.div_ceil(TILE), rows.div_ceil(TILE));
620 let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
621 label: Some("depth of field"),
622 timestamp_writes: timestamps,
623 });
624 pass.set_pipeline(&self.tile_pipeline);
625 pass.set_bind_group(0, &tiles, &[]);
626 pass.dispatch_workgroups(tiles_x, tiles_y, 1);
627 pass.set_pipeline(&self.dilate_pipeline);
628 pass.set_bind_group(0, &dilate, &[]);
629 pass.dispatch_workgroups(tiles_x.div_ceil(8), tiles_y.div_ceil(8), 1);
630 pass.set_pipeline(&self.gather_pipeline);
631 pass.set_bind_group(0, &gather, &[]);
632 pass.dispatch_workgroups(width.div_ceil(8), rows.div_ceil(8), 1);
633 &targets.output
634 }
635}
636
637#[cfg(test)]
638mod tests {
639 use super::*;
640 use crate::gpu::tests::{CAPACITY, VIEWPORTS, framed, half};
641
642 const TAN_HALF_WIDTH: f32 = 0.5;
643 const FRAME: u32 = 800;
644
645 fn lens() -> Lens {
646 Lens {
647 focal_m: 0.05,
648 sensor_m: Lens::sensor_for_view(0.05, 1.0, TAN_HALF_WIDTH),
649 fstop: 2.0,
650 aperture: 0.5,
651 focus_m: 1.0,
652 near_m: 1.0,
653 near_blur: 12.0,
654 far_m: 1.5,
655 far_blur: 24.0,
656 ground_m: None,
657 }
658 }
659
660 fn plain() -> Lens {
661 Lens {
662 near_blur: 0.0,
663 far_blur: 0.0,
664 ..lens()
665 }
666 }
667
668 #[test]
669 fn coc_matches_thin_lens_formula() {
670 let lens = plain();
671 for distance in [0.5_f32, 0.8, 1.5, 3.0] {
672 let blur = (lens.focal_m * lens.aperture / lens.fstop) * lens.focal_m
673 / (lens.focus_m - lens.focal_m)
674 * (distance - lens.focus_m)
675 / distance;
676 let expected =
677 (blur * FRAME as f32 / (2.0 * lens.sensor_m)).clamp(-MAX_RADIUS, MAX_RADIUS);
678 let got = lens.circle_of_confusion(distance, 0.3, FRAME);
679 assert!(
680 (got - expected).abs() < 1e-4,
681 "{distance}: {got} against {expected}"
682 );
683 }
684 }
685
686 #[test]
687 fn coc_matches_the_lens_disc() {
688 let lens = lens();
689 let radius = lens.focal_m / (2.0 * lens.fstop) * lens.aperture;
690 for (distance, height) in [(0.6_f32, 0.0_f32), (1.6, 0.4), (2.5, 0.4)] {
691 let open = lens.opening(distance, height);
692 let spread =
693 radius * open * (distance - lens.focus_m).abs() / (distance * lens.focus_m);
694 let expected = (spread * FRAME as f32 / (2.0 * TAN_HALF_WIDTH)).min(MAX_RADIUS);
695 let got = lens.circle_of_confusion(distance, height, FRAME).abs();
696 assert!(
697 (got - expected).abs() < 1e-3,
698 "{distance}: {got} against {expected}"
699 );
700 }
701 }
702
703 #[test]
704 fn near_blur_is_only_on_the_ground() {
705 let lens = Lens {
706 ground_m: Some(0.1),
707 ..lens()
708 };
709 assert!((lens.opening(0.5, 0.0) - (1.0 + 12.0 * 0.5 / 1.0)).abs() < 1e-5);
710 assert_eq!(lens.opening(0.5, 0.2), 1.0);
711 assert!((lens.opening(2.0, 0.2) - (1.0 + 24.0 * 0.5 / 1.5)).abs() < 1e-5);
712 }
713
714 #[test]
715 fn focus_plane_stays_sharp() {
716 assert_eq!(lens().circle_of_confusion(lens().focus_m, 0.0, FRAME), 0.0);
717 let (width, height) = (48, 32);
718 let image: Vec<_> = (0..width * height)
719 .map(|i| {
720 [
721 (i % 7) as f32 / 7.0,
722 (i % 5) as f32 / 5.0,
723 (i % 3) as f32,
724 1.0,
725 ]
726 })
727 .collect();
728 let depth = vec![lens().focus_m; width * height];
729 let heights = vec![0.0; width * height];
730 assert_eq!(
731 gather(&image, &depth, &heights, width, height, lens()),
732 image
733 );
734 }
735
736 fn edge(near: f32, far: f32, lens: Lens) -> (Vec<[f32; 4]>, usize) {
737 let (width, height) = (160, 48);
738 let image: Vec<_> = (0..width * height)
739 .map(|i| {
740 if i % width < 80 {
741 [1.0, 0.0, 0.0, 1.0]
742 } else {
743 [0.0, 0.0, 1.0, 1.0]
744 }
745 })
746 .collect();
747 let depth: Vec<_> = (0..width * height)
748 .map(|i| if i % width < 80 { near } else { far })
749 .collect();
750 let heights = vec![0.0; width * height];
751 (gather(&image, &depth, &heights, width, height, lens), width)
752 }
753
754 #[test]
755 fn sharp_foreground_takes_no_halo_from_a_blurred_background() {
756 let (out, width) = edge(lens().focus_m, 3.0, lens());
757 let row = 24 * width;
758 for x in 0..80 {
759 assert_eq!(out[row + x], [1.0, 0.0, 0.0, 1.0], "foreground at {x}");
760 }
761 for x in 80..160 {
762 assert!(
763 out[row + x][0] < 1e-6,
764 "background at {x} took {:?}",
765 out[row + x]
766 );
767 }
768 }
769
770 #[test]
771 fn blurred_foreground_spreads_over_a_sharp_background_without_a_halo() {
772 let wide = Lens {
773 aperture: 4.0,
774 ..lens()
775 };
776 let (out, width) = edge(0.5, lens().focus_m, wide);
777 let row = 24 * width;
778 let reach = wide.circle_of_confusion(0.6, 0.0, 160).abs();
779 assert!(reach > 4.0);
780 assert!(out[row + 81][0] > 0.2, "{:?}", out[row + 81]);
781 for x in 0..159 {
782 let here = out[row + x][0] + out[row + x][2];
783 assert!((here - 1.0).abs() < 1e-4, "energy at {x}: {here}");
784 assert!(
785 out[row + x + 1][0] <= out[row + x][0] + 1e-4,
786 "red rises at {x}"
787 );
788 }
789 assert_eq!(out[row + 159], [0.0, 0.0, 1.0, 1.0]);
790 }
791
792 #[test]
793 fn blur_keeps_surface_detail_above_the_ground() {
794 let lens = Lens {
795 ground_m: Some(0.1),
796 ..lens()
797 };
798 let (width, height) = (96, 64);
799 let image: Vec<_> = (0..width * height)
800 .map(|i| {
801 if (i % width) % 4 < 2 {
802 [1.0; 4]
803 } else {
804 [0.0, 0.0, 0.0, 1.0]
805 }
806 })
807 .collect();
808 let depth = vec![2.5; width * height];
809 let wall = vec![0.4; width * height];
810 assert_eq!(gather(&image, &depth, &wall, width, height, lens), image);
811 let floor = vec![0.0; width * height];
812 let blurred = gather(&image, &depth, &floor, width, height, lens);
813 let row: Vec<f32> = (16..80).map(|x| blurred[32 * width + x][0]).collect();
814 let (low, high) = row.iter().fold((f32::MAX, f32::MIN), |(low, high), &v| {
815 (low.min(v), high.max(v))
816 });
817 assert!(
818 high - low < 0.5,
819 "the floor kept its stripes: {low}..{high}"
820 );
821 }
822
823 #[test]
824 fn height_plane_finds_the_world_height() {
825 let (tan_x, tan_y, shift) = (0.5_f32, 0.3_f32, 0.1_f32);
826 let eye = [0.0_f32, 1.0, 2.0];
827 let view = [
828 [1.0, 0.0, 0.0, 0.0],
829 [0.0, 1.0, 0.0, 0.0],
830 [0.0, 0.0, 1.0, 0.0],
831 [-eye[0], -eye[1], -eye[2], 1.0],
832 ];
833 let projection = [
834 [1.0 / tan_x, 0.0, 0.0, 0.0],
835 [0.0, 1.0 / tan_y, 0.0, 0.0],
836 [0.0, -shift / tan_y, -1.0, -1.0],
837 [0.0, 0.0, -0.01, 0.0],
838 ];
839 let plane = height_plane(view, projection);
840 for (ndc, depth) in [([0.0_f32, -1.0_f32], 0.75_f32), ([0.5, 0.3], 1.2)] {
841 let ray_y = ndc[1] * tan_y - shift;
842 let expected = eye[1] + ray_y * depth;
843 let got =
844 plane[0] * ndc[0] * depth + plane[1] * ndc[1] * depth + plane[2] * depth + plane[3];
845 assert!((got - expected).abs() < 1e-5, "{got} against {expected}");
846 }
847 }
848
849 #[test]
850 fn shaders_parse() {
851 for part in [TILE_WGSL, DILATE_WGSL, GATHER_WGSL] {
852 naga::front::wgsl::parse_str(&shader(part)).unwrap();
853 }
854 }
855
856 const FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba16Float;
857
858 fn colour(x: u32, y: u32) -> [f32; 4] {
859 [
860 ((x * 7 + y * 3) % 17) as f32 / 17.0,
861 ((x / 6 + y / 4) % 5) as f32 / 4.0,
862 ((x + 2 * y) % 11) as f32 / 10.0,
863 1.0,
864 ]
865 }
866
867 fn distance(x: u32, y: u32) -> [f32; 4] {
868 let d = match (x / 13 + y / 9) % 3 {
869 0 => 0.6,
870 1 => lens().focus_m,
871 _ => 2.5,
872 };
873 [d, 0.0, 0.0, 1.0]
874 }
875
876 struct Frame {
877 source: pfx_gpu::OffscreenTarget,
878 depth: pfx_gpu::OffscreenTarget,
879 }
880
881 fn frame(gpu: &pfx_gpu::Gpu, size: (u32, u32), active: (u32, u32)) -> Frame {
882 let source = gpu.offscreen(size.0, size.1, FORMAT).unwrap();
883 gpu.upload_rgba16(
884 &source,
885 &framed(size, active, colour, [400.0, 0.0, 900.0, 1.0]),
886 )
887 .unwrap();
888 let depth = gpu.offscreen(size.0, size.1, FORMAT).unwrap();
889 gpu.upload_rgba16(
890 &depth,
891 &framed(size, active, distance, [0.2, 0.0, 0.0, 1.0]),
892 )
893 .unwrap();
894 Frame { source, depth }
895 }
896
897 fn render(gpu: &pfx_gpu::Gpu, dof: &GpuDof, frame: &Frame) -> Vec<u16> {
898 let mut encoder = gpu
899 .device
900 .create_command_encoder(&wgpu::CommandEncoderDescriptor {
901 label: Some("depth of field viewport test"),
902 });
903 dof.render(
904 &mut encoder,
905 &frame.source.view,
906 &frame.depth.view,
907 lens(),
908 [0.0, 1.0, 0.0, -0.5],
909 None,
910 );
911 gpu.queue.submit(Some(encoder.finish()));
912 let (width, height) = dof.capacity();
913 gpu.readback_rgba16(&pfx_gpu::OffscreenTarget {
914 texture: dof.targets._output.clone(),
915 view: dof.targets.output.clone(),
916 format: FORMAT,
917 width,
918 height,
919 })
920 .unwrap()
921 }
922
923 #[test]
924 #[ignore = "needs a GPU; run under pgpu"]
925 fn a_viewport_dof_equals_dof_at_the_viewport_size() {
926 let gpu = pollster::block_on(pfx_gpu::Gpu::headless()).unwrap();
927 let garbage = frame(&gpu, CAPACITY, (0, 0));
928 let mut dof = GpuDof::new(&gpu.device, CAPACITY.0, CAPACITY.1);
929 for viewport in VIEWPORTS {
930 dof.set_viewport(CAPACITY.0, CAPACITY.1).unwrap();
931 let before = render(&gpu, &dof, &garbage);
932 dof.set_viewport(viewport.0, viewport.1).unwrap();
933 assert_eq!(dof.viewport(), viewport);
934 let got = render(&gpu, &dof, &frame(&gpu, CAPACITY, viewport));
935 let reference = GpuDof::new(&gpu.device, viewport.0, viewport.1);
936 let want = render(&gpu, &reference, &frame(&gpu, viewport, viewport));
937 let mut blurred = 0;
938 for y in 0..CAPACITY.1 {
939 for x in 0..CAPACITY.0 {
940 let at = ((y * CAPACITY.0 + x) * 4) as usize;
941 if x >= viewport.0 || y >= viewport.1 {
942 assert_eq!(
943 got[at..at + 4],
944 before[at..at + 4],
945 "{viewport:?} wrote ({x},{y}) outside its viewport"
946 );
947 continue;
948 }
949 let there = ((y * viewport.0 + x) * 4) as usize;
950 assert_eq!(
951 got[at..at + 4],
952 want[there..there + 4],
953 "{viewport:?} ({x},{y}) must equal depth of field at its size exactly"
954 );
955 let sharp = colour(x, y).map(half);
956 if got[at..at + 4] != sharp {
957 blurred += 1;
958 }
959 }
960 }
961 assert!(blurred > 1000, "{viewport:?} blurred only {blurred} pixels");
962 }
963 }
964
965 #[test]
966 #[ignore = "needs a GPU; run under pgpu"]
967 fn a_full_viewport_dof_is_byte_identical() {
968 let gpu = pollster::block_on(pfx_gpu::Gpu::headless()).unwrap();
969 let whole = frame(&gpu, CAPACITY, CAPACITY);
970 let fresh = GpuDof::new(&gpu.device, CAPACITY.0, CAPACITY.1);
971 let want = render(&gpu, &fresh, &whole);
972 let mut dof = GpuDof::new(&gpu.device, CAPACITY.0, CAPACITY.1);
973 dof.set_viewport(VIEWPORTS[1].0, VIEWPORTS[1].1).unwrap();
974 render(&gpu, &dof, &frame(&gpu, CAPACITY, (0, 0)));
975 dof.set_viewport(CAPACITY.0, CAPACITY.1).unwrap();
976 assert!(render(&gpu, &dof, &whole) == want);
977 assert!(dof.set_viewport(0, 10).is_err());
978 assert!(dof.set_viewport(161, 10).is_err());
979 dof.set_viewport(50, 40).unwrap();
980 dof.resize(CAPACITY.0, CAPACITY.1);
981 assert_eq!(dof.viewport(), CAPACITY);
982 }
983}