dr-gpu: the merge pass — warp, accumulate, resolve, chunk by chunk
merge.wgsl warps one camera-space tile into one output chunk — output pixel to direction (the projection maths of dr_pano::projection, verbatim), direction to the frame's camera, camera to source pixel, bilinear by hand from four textureLoads because rgba32float is not filterable — and adds it into a storage-buffer accumulator weighted by its distance from the frame's edge. A resolve pass divides by the weights and packs sixteen-bit samples at the sensor's scale with a coverage bit. MergePass::merge drives it: bands of rows, chunks across a band, and for each chunk only the frames whose footprint meets it, each rendered as the source rectangle the chunk needs and nothing more. The working set is one chunk, one tile and one band (FR-MRG-11); the frame textures are the caller's to cache. Feathered, not seamed; gain a scalar per frame — the blend quality is panorama.md §10's step 5, after the path writes a file.
This commit is contained in:
@@ -7,6 +7,11 @@ license.workspace = true
|
|||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
dr-types.workspace = true
|
dr-types.workspace = true
|
||||||
|
# The merge's geometry (FR-MRG-10): rotations, the focal length and the
|
||||||
|
# projections, solved on proxies by dr-pano and consumed here per chunk. The
|
||||||
|
# geometry alone — no keypoint model, no runtime — which is what the
|
||||||
|
# workspace entry turns off.
|
||||||
|
dr-pano.workspace = true
|
||||||
dr-decode.workspace = true
|
dr-decode.workspace = true
|
||||||
dr-pipeline.workspace = true
|
dr-pipeline.workspace = true
|
||||||
# The watershed's pixel passes are here because they are shaders; everything
|
# The watershed's pixel passes are here because they are shaders; everything
|
||||||
|
|||||||
@@ -28,6 +28,7 @@ mod error;
|
|||||||
mod focus;
|
mod focus;
|
||||||
mod histogram;
|
mod histogram;
|
||||||
mod mask;
|
mod mask;
|
||||||
|
mod merge;
|
||||||
mod raw_histogram;
|
mod raw_histogram;
|
||||||
mod readback;
|
mod readback;
|
||||||
mod segment;
|
mod segment;
|
||||||
@@ -37,6 +38,7 @@ pub use adjust::AdjustPass;
|
|||||||
// unclipped, full internal precision (FR-DEV-2), and anyone reasoning about
|
// unclipped, full internal precision (FR-DEV-2), and anyone reasoning about
|
||||||
// VRAM at 24 MP needs to know what an intermediate costs.
|
// VRAM at 24 MP needs to know what an intermediate costs.
|
||||||
pub use demosaic::{DemosaicedImage, Demosaicer};
|
pub use demosaic::{DemosaicedImage, Demosaicer};
|
||||||
|
pub use merge::{Band, MergeFrame, MergeOutput, MergePass};
|
||||||
pub use detail::INTERMEDIATE_FORMAT as DETAIL_INTERMEDIATE_FORMAT;
|
pub use detail::INTERMEDIATE_FORMAT as DETAIL_INTERMEDIATE_FORMAT;
|
||||||
pub use error::GpuError;
|
pub use error::GpuError;
|
||||||
pub use focus::{FocusPeakPass, FocusPeaking, PeakColour, PeakSensitivity};
|
pub use focus::{FocusPeakPass, FocusPeaking, PeakColour, PeakSensitivity};
|
||||||
|
|||||||
@@ -0,0 +1,539 @@
|
|||||||
|
//! TRACES: FR-MRG-10 | FR-MRG-11
|
||||||
|
//! The merge: source frames warped into an output surface, chunk by chunk.
|
||||||
|
//!
|
||||||
|
//! The per-pixel half of a panorama (FR-MRG-10), on the GPU: the warp of a
|
||||||
|
//! source tile into an output chunk, the weighted accumulation across
|
||||||
|
//! frames, and the resolve to sixteen-bit samples. The geometry it is
|
||||||
|
//! given — rotations, focal length, projection — is `dr-pano`'s, solved on
|
||||||
|
//! proxies before any full-resolution pixel exists (panorama.md §5), and
|
||||||
|
//! that is what makes this simple: every output pixel's source coordinates
|
||||||
|
//! are a closed-form function, so a chunk can be produced from the source
|
||||||
|
//! tiles that project into it and nothing else.
|
||||||
|
//!
|
||||||
|
//! # The loop
|
||||||
|
//!
|
||||||
|
//! ```text
|
||||||
|
//! for each band of rows of the output:
|
||||||
|
//! for each chunk across the band:
|
||||||
|
//! zero the accumulator
|
||||||
|
//! for each frame whose footprint meets the chunk:
|
||||||
|
//! the source rectangle the chunk needs, from the geometry
|
||||||
|
//! render it camera-linear through the pipeline (the tile)
|
||||||
|
//! warp the tile into the chunk, accumulate ← GPU
|
||||||
|
//! resolve the chunk to u16 ← GPU
|
||||||
|
//! copy it into the band
|
||||||
|
//! hand the band to the writer (one DNG strip)
|
||||||
|
//! ```
|
||||||
|
//!
|
||||||
|
//! No stage holds the composite (FR-MRG-11): the working set is one
|
||||||
|
//! chunk's accumulator, one tile, one band of u16 rows. The frame textures
|
||||||
|
//! are the caller's to provide and cache — `source` is asked for frame `k`
|
||||||
|
//! as it is needed, and a caller short of memory may demosaic on demand.
|
||||||
|
//!
|
||||||
|
//! # What is not here yet
|
||||||
|
//!
|
||||||
|
//! A feathered blend, not seams and a Laplacian pyramid: the weight is the
|
||||||
|
//! distance to the frame's edge, which hides exposure steps and small
|
||||||
|
//! misalignments and does not hide parallax. Gain is a scalar per frame
|
||||||
|
//! the caller supplies. Both are panorama.md §10's step 5, after the path
|
||||||
|
//! writes a file end to end.
|
||||||
|
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use dr_pano::bundle::Cameras;
|
||||||
|
use dr_pano::projection::{Bounds, Projection};
|
||||||
|
use wgpu::util::DeviceExt;
|
||||||
|
|
||||||
|
use crate::readback::await_mapping;
|
||||||
|
use crate::{AdjustPass, DemosaicedImage, GpuContext, GpuError};
|
||||||
|
|
||||||
|
/// One frame's part in the merge.
|
||||||
|
pub struct MergeFrame {
|
||||||
|
/// The frame's edit, for its lens corrections — the only part of an
|
||||||
|
/// edit the camera-space tap uses (FR-MRG-2).
|
||||||
|
pub graph: Arc<dr_pipeline::EditGraph>,
|
||||||
|
/// Multiplies the frame's samples, to bring its exposure to the
|
||||||
|
/// reference frame's. 1.0 for no correction.
|
||||||
|
pub gain: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The output the merge produces.
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||||
|
pub struct MergeOutput {
|
||||||
|
pub projection: Projection,
|
||||||
|
/// The projection's scale in output pixels: the cylinder's radius, the
|
||||||
|
/// plane's distance. The source focal length at full resolution gives
|
||||||
|
/// output pixels the size of source pixels at the centre.
|
||||||
|
pub scale: f64,
|
||||||
|
/// The rectangle of the projection to produce, centred coordinates.
|
||||||
|
pub bounds: Bounds,
|
||||||
|
/// Pixels over which a frame's weight ramps up from its edge.
|
||||||
|
pub feather: f32,
|
||||||
|
/// Chunk size: the unit of GPU work and of memory.
|
||||||
|
pub chunk: (u32, u32),
|
||||||
|
/// Multiplies a normalised sample (1.0 = white) to the sensor's scale.
|
||||||
|
pub sample_scale: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl MergeOutput {
|
||||||
|
pub fn width(&self) -> u32 {
|
||||||
|
self.bounds.width().ceil().max(1.0) as u32
|
||||||
|
}
|
||||||
|
pub fn height(&self) -> u32 {
|
||||||
|
self.bounds.height().ceil().max(1.0) as u32
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A band of finished rows: `rows × width × 3` RGB `u16`, plus a coverage
|
||||||
|
/// mask (`true` where any frame reached the pixel).
|
||||||
|
pub struct Band<'a> {
|
||||||
|
pub first_row: u32,
|
||||||
|
pub rows: u32,
|
||||||
|
pub rgb: &'a [u16],
|
||||||
|
pub covered: &'a [bool],
|
||||||
|
}
|
||||||
|
|
||||||
|
#[repr(C)]
|
||||||
|
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
|
||||||
|
struct WarpParams {
|
||||||
|
chunk_origin: [f32; 2],
|
||||||
|
chunk_size: [u32; 2],
|
||||||
|
projection: u32,
|
||||||
|
proj_scale: f32,
|
||||||
|
focal: f32,
|
||||||
|
gain: f32,
|
||||||
|
r0: [f32; 4],
|
||||||
|
r1: [f32; 4],
|
||||||
|
r2: [f32; 4],
|
||||||
|
frame_size: [f32; 2],
|
||||||
|
tile_origin: [f32; 2],
|
||||||
|
tile_size: [u32; 2],
|
||||||
|
feather: f32,
|
||||||
|
_pad: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[repr(C)]
|
||||||
|
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
|
||||||
|
struct ResolveParams {
|
||||||
|
chunk_size: [u32; 2],
|
||||||
|
scale: f32,
|
||||||
|
_pad: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The two pipelines and the chunk buffers.
|
||||||
|
pub struct MergePass {
|
||||||
|
ctx: GpuContext,
|
||||||
|
warp: wgpu::ComputePipeline,
|
||||||
|
warp_layout: wgpu::BindGroupLayout,
|
||||||
|
resolve: wgpu::ComputePipeline,
|
||||||
|
resolve_layout: wgpu::BindGroupLayout,
|
||||||
|
/// Accumulator and packed output for the current chunk size.
|
||||||
|
buffers: Option<(wgpu::Buffer, wgpu::Buffer, wgpu::Buffer, (u32, u32))>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl MergePass {
|
||||||
|
pub fn new(ctx: &GpuContext) -> Result<Self, GpuError> {
|
||||||
|
let module = ctx
|
||||||
|
.device
|
||||||
|
.create_shader_module(wgpu::ShaderModuleDescriptor {
|
||||||
|
label: Some("merge"),
|
||||||
|
source: wgpu::ShaderSource::Wgsl(include_str!("shaders/merge.wgsl").into()),
|
||||||
|
});
|
||||||
|
let uniform = |binding| wgpu::BindGroupLayoutEntry {
|
||||||
|
binding,
|
||||||
|
visibility: wgpu::ShaderStages::COMPUTE,
|
||||||
|
ty: wgpu::BindingType::Buffer {
|
||||||
|
ty: wgpu::BufferBindingType::Uniform,
|
||||||
|
has_dynamic_offset: false,
|
||||||
|
min_binding_size: None,
|
||||||
|
},
|
||||||
|
count: None,
|
||||||
|
};
|
||||||
|
let storage = |binding, read_only| wgpu::BindGroupLayoutEntry {
|
||||||
|
binding,
|
||||||
|
visibility: wgpu::ShaderStages::COMPUTE,
|
||||||
|
ty: wgpu::BindingType::Buffer {
|
||||||
|
ty: wgpu::BufferBindingType::Storage { read_only },
|
||||||
|
has_dynamic_offset: false,
|
||||||
|
min_binding_size: None,
|
||||||
|
},
|
||||||
|
count: None,
|
||||||
|
};
|
||||||
|
let warp_layout = ctx
|
||||||
|
.device
|
||||||
|
.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
|
||||||
|
label: Some("merge-warp-bgl"),
|
||||||
|
entries: &[
|
||||||
|
uniform(0),
|
||||||
|
wgpu::BindGroupLayoutEntry {
|
||||||
|
binding: 1,
|
||||||
|
visibility: wgpu::ShaderStages::COMPUTE,
|
||||||
|
ty: wgpu::BindingType::Texture {
|
||||||
|
// Unfilterable: rgba32float, loaded by hand.
|
||||||
|
sample_type: wgpu::TextureSampleType::Float { filterable: false },
|
||||||
|
view_dimension: wgpu::TextureViewDimension::D2,
|
||||||
|
multisampled: false,
|
||||||
|
},
|
||||||
|
count: None,
|
||||||
|
},
|
||||||
|
storage(2, false),
|
||||||
|
],
|
||||||
|
});
|
||||||
|
let resolve_layout = ctx
|
||||||
|
.device
|
||||||
|
.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
|
||||||
|
label: Some("merge-resolve-bgl"),
|
||||||
|
entries: &[uniform(0), storage(1, true), storage(2, false)],
|
||||||
|
});
|
||||||
|
let pipeline = |name: &str, layout: &wgpu::BindGroupLayout| {
|
||||||
|
let pl = ctx
|
||||||
|
.device
|
||||||
|
.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
|
||||||
|
label: Some(name),
|
||||||
|
bind_group_layouts: &[Some(layout)],
|
||||||
|
immediate_size: 0,
|
||||||
|
});
|
||||||
|
ctx.device
|
||||||
|
.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
|
||||||
|
label: Some(name),
|
||||||
|
layout: Some(&pl),
|
||||||
|
module: &module,
|
||||||
|
entry_point: Some(name),
|
||||||
|
compilation_options: Default::default(),
|
||||||
|
cache: None,
|
||||||
|
})
|
||||||
|
};
|
||||||
|
Ok(MergePass {
|
||||||
|
ctx: ctx.clone(),
|
||||||
|
warp: pipeline("warp", &warp_layout),
|
||||||
|
warp_layout,
|
||||||
|
resolve: pipeline("resolve", &resolve_layout),
|
||||||
|
resolve_layout,
|
||||||
|
buffers: None,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Allocate the chunk buffers for this size if the last ones differ.
|
||||||
|
fn ensure_buffers(&mut self, chunk: (u32, u32)) {
|
||||||
|
if self.buffers.as_ref().is_none_or(|b| b.3 != chunk) {
|
||||||
|
let n = u64::from(chunk.0) * u64::from(chunk.1);
|
||||||
|
let acc = self.ctx.device.create_buffer(&wgpu::BufferDescriptor {
|
||||||
|
label: Some("merge-acc"),
|
||||||
|
size: n * 16,
|
||||||
|
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
|
||||||
|
mapped_at_creation: false,
|
||||||
|
});
|
||||||
|
let out = self.ctx.device.create_buffer(&wgpu::BufferDescriptor {
|
||||||
|
label: Some("merge-out"),
|
||||||
|
size: n * 8,
|
||||||
|
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
|
||||||
|
mapped_at_creation: false,
|
||||||
|
});
|
||||||
|
let read = self.ctx.device.create_buffer(&wgpu::BufferDescriptor {
|
||||||
|
label: Some("merge-read"),
|
||||||
|
size: n * 8,
|
||||||
|
usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
|
||||||
|
mapped_at_creation: false,
|
||||||
|
});
|
||||||
|
self.buffers = Some((acc, out, read, chunk));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn chunk_buffers(&self) -> (&wgpu::Buffer, &wgpu::Buffer, &wgpu::Buffer) {
|
||||||
|
let b = self.buffers.as_ref().expect("ensured by the caller");
|
||||||
|
(&b.0, &b.1, &b.2)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Produce the whole output, band by band, handing each finished band
|
||||||
|
/// to `sink`.
|
||||||
|
///
|
||||||
|
/// `cameras` are in **full-resolution source pixels** (`frame_size`),
|
||||||
|
/// with frame `k` corresponding to `frames[k]` and `source(k)`. `source`
|
||||||
|
/// supplies the demosaiced frame on demand and may cache as it sees fit.
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
|
pub fn merge<S, F>(
|
||||||
|
&mut self,
|
||||||
|
adjust: &mut AdjustPass,
|
||||||
|
frames: &[MergeFrame],
|
||||||
|
cameras: &Cameras,
|
||||||
|
frame_size: (u32, u32),
|
||||||
|
output: &MergeOutput,
|
||||||
|
mut source: S,
|
||||||
|
mut sink: F,
|
||||||
|
mut cancelled: impl FnMut() -> bool,
|
||||||
|
) -> Result<(), GpuError>
|
||||||
|
where
|
||||||
|
S: FnMut(usize) -> Result<Arc<DemosaicedImage>, GpuError>,
|
||||||
|
F: FnMut(Band<'_>) -> Result<(), GpuError>,
|
||||||
|
{
|
||||||
|
let (out_w, out_h) = (output.width(), output.height());
|
||||||
|
let (cw, ch) = (output.chunk.0.max(8), output.chunk.1.max(8));
|
||||||
|
let (fw, fh) = (frame_size.0 as f64, frame_size.1 as f64);
|
||||||
|
|
||||||
|
let mut band_rgb = vec![0u16; (out_w * ch * 3) as usize];
|
||||||
|
let mut band_cov = vec![false; (out_w * ch) as usize];
|
||||||
|
let mut chunk_px: Vec<u32> = Vec::new();
|
||||||
|
|
||||||
|
let mut y = 0u32;
|
||||||
|
while y < out_h {
|
||||||
|
let rows = ch.min(out_h - y);
|
||||||
|
band_rgb.iter_mut().for_each(|v| *v = 0);
|
||||||
|
band_cov.iter_mut().for_each(|v| *v = false);
|
||||||
|
|
||||||
|
let mut x = 0u32;
|
||||||
|
while x < out_w {
|
||||||
|
if cancelled() {
|
||||||
|
return Err(GpuError::Readback("merge cancelled".into()));
|
||||||
|
}
|
||||||
|
let cols = cw.min(out_w - x);
|
||||||
|
let origin = (
|
||||||
|
output.bounds.min_u + f64::from(x),
|
||||||
|
output.bounds.min_v + f64::from(y),
|
||||||
|
);
|
||||||
|
self.zero_accumulator((cols, rows));
|
||||||
|
|
||||||
|
for (k, frame) in frames.iter().enumerate() {
|
||||||
|
let Some(rect) = source_rect(
|
||||||
|
output.projection,
|
||||||
|
output.scale,
|
||||||
|
cameras,
|
||||||
|
k,
|
||||||
|
origin,
|
||||||
|
(cols, rows),
|
||||||
|
(fw, fh),
|
||||||
|
) else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
let image = source(k)?;
|
||||||
|
// The tile: that rectangle of the frame, camera-linear,
|
||||||
|
// at 1:1.
|
||||||
|
let view = dr_pipeline::CropRect {
|
||||||
|
x: (rect.0 as f32) / fw as f32,
|
||||||
|
y: (rect.1 as f32) / fh as f32,
|
||||||
|
width: (rect.2 as f32) / fw as f32,
|
||||||
|
height: (rect.3 as f32) / fh as f32,
|
||||||
|
};
|
||||||
|
let shader = frame.graph.compose_camera_linear(view);
|
||||||
|
let tile = adjust.render_camera_linear(&image, &shader, rect.2, rect.3)?;
|
||||||
|
let r = cameras.rotations[k].transpose();
|
||||||
|
let params = WarpParams {
|
||||||
|
chunk_origin: [origin.0 as f32, origin.1 as f32],
|
||||||
|
chunk_size: [cols, rows],
|
||||||
|
projection: match output.projection {
|
||||||
|
Projection::Perspective => 0,
|
||||||
|
Projection::Cylindrical => 1,
|
||||||
|
Projection::Spherical => 2,
|
||||||
|
},
|
||||||
|
proj_scale: output.scale as f32,
|
||||||
|
focal: cameras.focal as f32,
|
||||||
|
gain: frame.gain,
|
||||||
|
r0: [r.0[0][0] as f32, r.0[0][1] as f32, r.0[0][2] as f32, 0.0],
|
||||||
|
r1: [r.0[1][0] as f32, r.0[1][1] as f32, r.0[1][2] as f32, 0.0],
|
||||||
|
r2: [r.0[2][0] as f32, r.0[2][1] as f32, r.0[2][2] as f32, 0.0],
|
||||||
|
frame_size: [fw as f32, fh as f32],
|
||||||
|
tile_origin: [rect.0 as f32, rect.1 as f32],
|
||||||
|
tile_size: [rect.2, rect.3],
|
||||||
|
feather: output.feather,
|
||||||
|
_pad: 0.0,
|
||||||
|
};
|
||||||
|
self.accumulate(¶ms, tile);
|
||||||
|
}
|
||||||
|
|
||||||
|
self.resolve_chunk((cols, rows), output.sample_scale, &mut chunk_px)?;
|
||||||
|
// Into the band.
|
||||||
|
for row in 0..rows as usize {
|
||||||
|
for col in 0..cols as usize {
|
||||||
|
let px = chunk_px[(row * cols as usize + col) * 2..][..2].to_vec();
|
||||||
|
let i = row * out_w as usize + (x as usize + col);
|
||||||
|
band_rgb[i * 3] = (px[0] & 0xFFFF) as u16;
|
||||||
|
band_rgb[i * 3 + 1] = (px[0] >> 16) as u16;
|
||||||
|
band_rgb[i * 3 + 2] = (px[1] & 0xFFFF) as u16;
|
||||||
|
band_cov[i] = (px[1] >> 16) != 0;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
x += cols;
|
||||||
|
}
|
||||||
|
|
||||||
|
sink(Band {
|
||||||
|
first_row: y,
|
||||||
|
rows,
|
||||||
|
rgb: &band_rgb[..(out_w * rows * 3) as usize],
|
||||||
|
covered: &band_cov[..(out_w * rows) as usize],
|
||||||
|
})?;
|
||||||
|
y += rows;
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn zero_accumulator(&mut self, chunk: (u32, u32)) {
|
||||||
|
self.ensure_buffers(chunk);
|
||||||
|
let (acc, _, _) = self.chunk_buffers();
|
||||||
|
let n = u64::from(chunk.0) * u64::from(chunk.1) * 16;
|
||||||
|
let mut enc = self.ctx.device.create_command_encoder(&Default::default());
|
||||||
|
enc.clear_buffer(acc, 0, Some(n));
|
||||||
|
self.ctx.queue.submit(Some(enc.finish()));
|
||||||
|
}
|
||||||
|
|
||||||
|
fn accumulate(&mut self, params: &WarpParams, tile: &wgpu::Texture) {
|
||||||
|
let chunk = (params.chunk_size[0], params.chunk_size[1]);
|
||||||
|
let uniforms = self
|
||||||
|
.ctx
|
||||||
|
.device
|
||||||
|
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
|
||||||
|
label: Some("merge-warp-params"),
|
||||||
|
contents: bytemuck::bytes_of(params),
|
||||||
|
usage: wgpu::BufferUsages::UNIFORM,
|
||||||
|
});
|
||||||
|
let view = tile.create_view(&Default::default());
|
||||||
|
self.ensure_buffers(chunk);
|
||||||
|
let (acc, _, _) = self.chunk_buffers();
|
||||||
|
let bind = self
|
||||||
|
.ctx
|
||||||
|
.device
|
||||||
|
.create_bind_group(&wgpu::BindGroupDescriptor {
|
||||||
|
label: Some("merge-warp-bg"),
|
||||||
|
layout: &self.warp_layout,
|
||||||
|
entries: &[
|
||||||
|
wgpu::BindGroupEntry {
|
||||||
|
binding: 0,
|
||||||
|
resource: uniforms.as_entire_binding(),
|
||||||
|
},
|
||||||
|
wgpu::BindGroupEntry {
|
||||||
|
binding: 1,
|
||||||
|
resource: wgpu::BindingResource::TextureView(&view),
|
||||||
|
},
|
||||||
|
wgpu::BindGroupEntry {
|
||||||
|
binding: 2,
|
||||||
|
resource: acc.as_entire_binding(),
|
||||||
|
},
|
||||||
|
],
|
||||||
|
});
|
||||||
|
let mut enc = self.ctx.device.create_command_encoder(&Default::default());
|
||||||
|
{
|
||||||
|
let mut pass = enc.begin_compute_pass(&Default::default());
|
||||||
|
pass.set_pipeline(&self.warp);
|
||||||
|
pass.set_bind_group(0, &bind, &[]);
|
||||||
|
pass.dispatch_workgroups(chunk.0.div_ceil(8), chunk.1.div_ceil(8), 1);
|
||||||
|
}
|
||||||
|
self.ctx.queue.submit(Some(enc.finish()));
|
||||||
|
}
|
||||||
|
|
||||||
|
fn resolve_chunk(&mut self, chunk: (u32, u32), scale: f32, out: &mut Vec<u32>) -> Result<(), GpuError> {
|
||||||
|
let params = ResolveParams {
|
||||||
|
chunk_size: [chunk.0, chunk.1],
|
||||||
|
scale,
|
||||||
|
_pad: 0.0,
|
||||||
|
};
|
||||||
|
let uniforms = self
|
||||||
|
.ctx
|
||||||
|
.device
|
||||||
|
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
|
||||||
|
label: Some("merge-resolve-params"),
|
||||||
|
contents: bytemuck::bytes_of(¶ms),
|
||||||
|
usage: wgpu::BufferUsages::UNIFORM,
|
||||||
|
});
|
||||||
|
let n = u64::from(chunk.0) * u64::from(chunk.1);
|
||||||
|
self.ensure_buffers(chunk);
|
||||||
|
let (acc, packed, read) = self.chunk_buffers();
|
||||||
|
let bind = self
|
||||||
|
.ctx
|
||||||
|
.device
|
||||||
|
.create_bind_group(&wgpu::BindGroupDescriptor {
|
||||||
|
label: Some("merge-resolve-bg"),
|
||||||
|
layout: &self.resolve_layout,
|
||||||
|
entries: &[
|
||||||
|
wgpu::BindGroupEntry {
|
||||||
|
binding: 0,
|
||||||
|
resource: uniforms.as_entire_binding(),
|
||||||
|
},
|
||||||
|
wgpu::BindGroupEntry {
|
||||||
|
binding: 1,
|
||||||
|
resource: acc.as_entire_binding(),
|
||||||
|
},
|
||||||
|
wgpu::BindGroupEntry {
|
||||||
|
binding: 2,
|
||||||
|
resource: packed.as_entire_binding(),
|
||||||
|
},
|
||||||
|
],
|
||||||
|
});
|
||||||
|
let mut enc = self.ctx.device.create_command_encoder(&Default::default());
|
||||||
|
{
|
||||||
|
let mut pass = enc.begin_compute_pass(&Default::default());
|
||||||
|
pass.set_pipeline(&self.resolve);
|
||||||
|
pass.set_bind_group(0, &bind, &[]);
|
||||||
|
pass.dispatch_workgroups(chunk.0.div_ceil(8), chunk.1.div_ceil(8), 1);
|
||||||
|
}
|
||||||
|
enc.copy_buffer_to_buffer(packed, 0, read, 0, n * 8);
|
||||||
|
self.ctx.queue.submit(Some(enc.finish()));
|
||||||
|
|
||||||
|
let slice = read.slice(..n * 8);
|
||||||
|
let (tx, rx) = std::sync::mpsc::channel();
|
||||||
|
slice.map_async(wgpu::MapMode::Read, move |r| {
|
||||||
|
let _ = tx.send(r);
|
||||||
|
});
|
||||||
|
await_mapping(&self.ctx, &rx)?;
|
||||||
|
{
|
||||||
|
let data = slice.get_mapped_range();
|
||||||
|
out.clear();
|
||||||
|
out.extend_from_slice(bytemuck::cast_slice::<u8, u32>(&data));
|
||||||
|
}
|
||||||
|
read.unmap();
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The rectangle of frame `k` (x, y, w, h in source pixels) a chunk reads,
|
||||||
|
/// or `None` if the chunk sees nothing of the frame.
|
||||||
|
///
|
||||||
|
/// Walks the chunk's border, projects each point into the frame, and takes
|
||||||
|
/// the bounding box with a two-pixel margin for the bilinear fetch. The
|
||||||
|
/// border rather than the corners because under a cylinder or sphere the
|
||||||
|
/// extreme of a footprint is not at a corner.
|
||||||
|
fn source_rect(
|
||||||
|
projection: Projection,
|
||||||
|
scale: f64,
|
||||||
|
cameras: &Cameras,
|
||||||
|
k: usize,
|
||||||
|
origin: (f64, f64),
|
||||||
|
size: (u32, u32),
|
||||||
|
frame: (f64, f64),
|
||||||
|
) -> Option<(u32, u32, u32, u32)> {
|
||||||
|
let (w, h) = (f64::from(size.0), f64::from(size.1));
|
||||||
|
let steps = 16;
|
||||||
|
let mut min = (f64::MAX, f64::MAX);
|
||||||
|
let mut max = (f64::MIN, f64::MIN);
|
||||||
|
let mut any = false;
|
||||||
|
let mut visit = |u: f64, v: f64| {
|
||||||
|
let d = projection.to_direction(scale, u, v);
|
||||||
|
if let Some((x, y)) = cameras.project(k, d) {
|
||||||
|
let (x, y) = (x + frame.0 / 2.0, y + frame.1 / 2.0);
|
||||||
|
min = (min.0.min(x), min.1.min(y));
|
||||||
|
max = (max.0.max(x), max.1.max(y));
|
||||||
|
any = true;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
for s in 0..=steps {
|
||||||
|
let t = f64::from(s) / f64::from(steps);
|
||||||
|
visit(origin.0 + w * t, origin.1);
|
||||||
|
visit(origin.0 + w * t, origin.1 + h);
|
||||||
|
visit(origin.0, origin.1 + h * t);
|
||||||
|
visit(origin.0 + w, origin.1 + h * t);
|
||||||
|
}
|
||||||
|
// The interior too, coarsely: a chunk can contain a frame entirely.
|
||||||
|
for i in 1..4 {
|
||||||
|
for j in 1..4 {
|
||||||
|
visit(origin.0 + w * f64::from(i) / 4.0, origin.1 + h * f64::from(j) / 4.0);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !any {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let x0 = (min.0.floor() - 2.0).max(0.0);
|
||||||
|
let y0 = (min.1.floor() - 2.0).max(0.0);
|
||||||
|
let x1 = (max.0.ceil() + 2.0).min(frame.0);
|
||||||
|
let y1 = (max.1.ceil() + 2.0).min(frame.1);
|
||||||
|
if x1 <= x0 || y1 <= y0 {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
Some((x0 as u32, y0 as u32, (x1 - x0) as u32, (y1 - y0) as u32))
|
||||||
|
}
|
||||||
@@ -0,0 +1,152 @@
|
|||||||
|
// TRACES: FR-MRG-10 | FR-MRG-11
|
||||||
|
// The merge: one source tile warped into one output chunk, accumulated.
|
||||||
|
//
|
||||||
|
// Two entry points. `warp` runs once per (chunk, frame): for every chunk
|
||||||
|
// pixel it asks which direction that pixel looks along, turns the
|
||||||
|
// direction into the frame's camera, projects it to a source pixel, and
|
||||||
|
// if that pixel is inside the tile that was rendered for this chunk,
|
||||||
|
// samples it and adds it — weighted by its distance from the frame's edge
|
||||||
|
// — into the accumulator. `resolve` runs once per chunk after every frame
|
||||||
|
// has been added: divides the sums by the weights and packs the result as
|
||||||
|
// sixteen-bit samples at the sensor's scale (FR-MRG-3).
|
||||||
|
//
|
||||||
|
// The accumulator is a buffer and not a storage texture, because WebGPU
|
||||||
|
// allows a read-write storage texture only in the 32-bit single-channel
|
||||||
|
// formats, and this wants four channels. The tile is sampled by hand from
|
||||||
|
// four `textureLoad`s rather than through a sampler, because `rgba32float`
|
||||||
|
// is not filterable without an optional feature, and the tile is
|
||||||
|
// `rgba32float` on purpose (panorama.md §5.1).
|
||||||
|
//
|
||||||
|
// The projection maths is `dr_pano::projection` verbatim; the two must
|
||||||
|
// agree, and a golden test compares them.
|
||||||
|
|
||||||
|
struct Params {
|
||||||
|
// Where the chunk's pixel (0, 0) sits in centred output coordinates,
|
||||||
|
// and the chunk's size.
|
||||||
|
chunk_origin: vec2<f32>,
|
||||||
|
chunk_size: vec2<u32>,
|
||||||
|
// 0 perspective, 1 cylindrical, 2 spherical; and the projection's
|
||||||
|
// scale (the cylinder's radius, the sphere's, the plane's distance) in
|
||||||
|
// output pixels.
|
||||||
|
projection: u32,
|
||||||
|
proj_scale: f32,
|
||||||
|
// The frame's focal length in source pixels, and the gain the frame's
|
||||||
|
// exposure is corrected by.
|
||||||
|
focal: f32,
|
||||||
|
gain: f32,
|
||||||
|
// World → this frame's camera: the transpose of its rotation, one row
|
||||||
|
// per vec4 (padded).
|
||||||
|
r0: vec4<f32>,
|
||||||
|
r1: vec4<f32>,
|
||||||
|
r2: vec4<f32>,
|
||||||
|
// The full frame's size in source pixels (for the edge weight), the
|
||||||
|
// tile's origin within the frame, and the tile's size.
|
||||||
|
frame_size: vec2<f32>,
|
||||||
|
tile_origin: vec2<f32>,
|
||||||
|
tile_size: vec2<u32>,
|
||||||
|
// Pixels over which the weight ramps from the edge to full.
|
||||||
|
feather: f32,
|
||||||
|
_pad: f32,
|
||||||
|
};
|
||||||
|
|
||||||
|
@group(0) @binding(0) var<uniform> p: Params;
|
||||||
|
@group(0) @binding(1) var tile: texture_2d<f32>;
|
||||||
|
// rgb·w summed, then w: four floats per chunk pixel.
|
||||||
|
@group(0) @binding(2) var<storage, read_write> acc: array<vec4<f32>>;
|
||||||
|
|
||||||
|
fn to_direction(u: f32, v: f32) -> vec3<f32> {
|
||||||
|
let s = p.proj_scale;
|
||||||
|
if (p.projection == 0u) {
|
||||||
|
return normalize(vec3<f32>(u, v, s));
|
||||||
|
}
|
||||||
|
if (p.projection == 1u) {
|
||||||
|
let theta = u / s;
|
||||||
|
return normalize(vec3<f32>(sin(theta), v / s, cos(theta)));
|
||||||
|
}
|
||||||
|
let theta = u / s;
|
||||||
|
let phi = v / s;
|
||||||
|
return vec3<f32>(sin(theta) * cos(phi), sin(phi), cos(theta) * cos(phi));
|
||||||
|
}
|
||||||
|
|
||||||
|
fn load(x: i32, y: i32) -> vec3<f32> {
|
||||||
|
return textureLoad(tile, vec2<i32>(x, y), 0).rgb;
|
||||||
|
}
|
||||||
|
|
||||||
|
@compute @workgroup_size(8, 8, 1)
|
||||||
|
fn warp(@builtin(global_invocation_id) gid: vec3<u32>) {
|
||||||
|
if (gid.x >= p.chunk_size.x || gid.y >= p.chunk_size.y) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let u = p.chunk_origin.x + f32(gid.x) + 0.5;
|
||||||
|
let v = p.chunk_origin.y + f32(gid.y) + 0.5;
|
||||||
|
let d = to_direction(u, v);
|
||||||
|
let c = vec3<f32>(dot(p.r0.xyz, d), dot(p.r1.xyz, d), dot(p.r2.xyz, d));
|
||||||
|
if (c.z <= 1e-6) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
// Source pixel, in the full frame, with the principal point at its
|
||||||
|
// centre. `- 0.5` puts pixel centres on integer coordinates for the
|
||||||
|
// bilinear fetch below.
|
||||||
|
let sx = p.focal * c.x / c.z + p.frame_size.x * 0.5 - 0.5;
|
||||||
|
let sy = p.focal * c.y / c.z + p.frame_size.y * 0.5 - 0.5;
|
||||||
|
// Weight: distance to the nearest frame edge, in pixels, over the
|
||||||
|
// feather. Zero outside the frame.
|
||||||
|
let edge = min(min(sx, p.frame_size.x - 1.0 - sx), min(sy, p.frame_size.y - 1.0 - sy));
|
||||||
|
if (edge <= 0.0) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let w = clamp(edge / max(p.feather, 1.0), 0.0, 1.0);
|
||||||
|
// Into the tile.
|
||||||
|
let tx = sx - p.tile_origin.x;
|
||||||
|
let ty = sy - p.tile_origin.y;
|
||||||
|
let tw = f32(p.tile_size.x);
|
||||||
|
let th = f32(p.tile_size.y);
|
||||||
|
if (tx < 0.0 || ty < 0.0 || tx > tw - 1.0 || ty > th - 1.0) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let x0 = i32(floor(tx));
|
||||||
|
let y0 = i32(floor(ty));
|
||||||
|
let x1 = min(x0 + 1, i32(p.tile_size.x) - 1);
|
||||||
|
let y1 = min(y0 + 1, i32(p.tile_size.y) - 1);
|
||||||
|
let fx = tx - f32(x0);
|
||||||
|
let fy = ty - f32(y0);
|
||||||
|
let top = mix(load(x0, y0), load(x1, y0), fx);
|
||||||
|
let bot = mix(load(x0, y1), load(x1, y1), fx);
|
||||||
|
let rgb = mix(top, bot, fy) * p.gain;
|
||||||
|
let i = gid.y * p.chunk_size.x + gid.x;
|
||||||
|
acc[i] = acc[i] + vec4<f32>(rgb * w, w);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Resolve: the accumulated chunk to sixteen-bit samples.
|
||||||
|
struct ResolveParams {
|
||||||
|
chunk_size: vec2<u32>,
|
||||||
|
// Multiplies a normalised value (1.0 = the sensor's white) back to the
|
||||||
|
// sensor's scale: the source's white minus its black (FR-MRG-3).
|
||||||
|
scale: f32,
|
||||||
|
_pad: f32,
|
||||||
|
};
|
||||||
|
|
||||||
|
@group(0) @binding(0) var<uniform> rp: ResolveParams;
|
||||||
|
@group(0) @binding(1) var<storage, read> racc: array<vec4<f32>>;
|
||||||
|
// Two u32 per pixel: (r | g << 16), (b | coverage << 16). Coverage is
|
||||||
|
// 65535 where any frame reached the pixel and 0 where none did, so the
|
||||||
|
// CPU can tell an empty pixel from a black one.
|
||||||
|
@group(0) @binding(2) var<storage, read_write> out: array<vec2<u32>>;
|
||||||
|
|
||||||
|
@compute @workgroup_size(8, 8, 1)
|
||||||
|
fn resolve(@builtin(global_invocation_id) gid: vec3<u32>) {
|
||||||
|
if (gid.x >= rp.chunk_size.x || gid.y >= rp.chunk_size.y) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let i = gid.y * rp.chunk_size.x + gid.x;
|
||||||
|
let a = racc[i];
|
||||||
|
if (a.w <= 0.0) {
|
||||||
|
out[i] = vec2<u32>(0u, 0u);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let rgb = clamp(a.rgb / a.w * rp.scale, vec3<f32>(0.0), vec3<f32>(65535.0));
|
||||||
|
let r = u32(round(rgb.r));
|
||||||
|
let g = u32(round(rgb.g));
|
||||||
|
let b = u32(round(rgb.b));
|
||||||
|
out[i] = vec2<u32>(r | (g << 16u), b | (65535u << 16u));
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user