dtourolle
42d11d919b
The dr-face comparison is master's: a negated partial-order test on the eye box's width, rewritten as the two conditions it meant.
548 lines
21 KiB
Rust
548 lines
21 KiB
Rust
//! 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))
|
||
}
|