Files
DarkRoom/core/dr-gpu/src/merge.rs
T
dtourolle 6f33517b35 Cut panorama overlaps along seams instead of averaging them
The merge weighted every overlap pixel by its distance from each frame's
edge, a 200 px linear cross-fade. Anything the frames disagreed on —
parallax in the near foreground, grass in the wind, a walker — came out
twice at half strength: a soft double edge at 1:1.

dr_pano::seam picks, per output texel at proxy resolution, which frame a
pixel comes from. Where a new frame overlaps the composite the cost is the
gain-corrected difference plus local detail plus nearness to either
frame's edge, taken as the worst over a small window, and the cut is a
dynamic-programming path across the overlap. merge.wgsl weights each frame
by its tent-filtered share of that map, a 64 px blend that follows the
seam, with the edge feather kept as the fallback. The page's preview uses
the same map, and examples/merge.rs takes --feather-only for comparison.
2026-09-30 21:57:44 -04:00

660 lines
25 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! 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.
//!
//! # The blend
//!
//! With a seam map (`dr_pano::seam`), a frame's weight at a pixel is its
//! share of the map about that pixel — whole on its own side of a seam,
//! nothing on the other, and a ramp across a window `seam_blend` pixels
//! wide that follows the seam. Without one, or where the map has nothing
//! to say, the weight is the distance to the frame's edge over `feather`,
//! which hides exposure steps and does not hide parallax: the average draws
//! anything the frames disagree on twice.
//!
//! # What is not here yet
//!
//! A Laplacian pyramid, which would let the seam's blend be narrow for
//! detail and wide for exposure at once. Gain is a scalar per frame the
//! caller supplies.
use std::sync::Arc;
use dr_pano::bundle::Cameras;
use dr_pano::projection::{Bounds, Projection};
use dr_pano::seam::SeamMap;
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, 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,
/// Which frame each part of the output is taken from, laid out at the
/// proxies' scale; `None` for the feathered average everywhere.
pub seams: Option<Arc<SeamMap>>,
/// The width, in output pixels, of the blend across a seam.
pub seam_blend: 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,
/// The white balance the composite will be developed with — the
/// inverse of its `AsShotNeutral` — so that a blown sample can be
/// written as the camera value that balance calls grey.
pub balance: [f32; 3],
}
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,
clip_onset: f32,
balance: [f32; 4],
seam_origin: [f32; 2],
seam_size: [u32; 2],
seam_px: f32,
seam_radius: f32,
frame_index: u32,
seam_on: u32,
}
#[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),
wgpu::BindGroupLayoutEntry {
binding: 3,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Texture {
sample_type: wgpu::TextureSampleType::Uint,
view_dimension: wgpu::TextureViewDimension::D2,
multisampled: false,
},
count: None,
},
],
});
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();
// The seam map, once for the whole output, and where it sits in
// this output's coordinates. A one-texel stand-in when there is
// none, because the binding is not optional.
let (seam_tex, seam_origin, seam_px, seam_radius, seam_size) = match &output.seams {
Some(m) => {
let ((ou, ov), px) = m.at_scale(output.scale);
let radius = m.blend_radius(output.scale, f64::from(output.seam_blend));
(
self.label_texture(m.width as u32, m.height as u32, &m.labels),
[ou as f32, ov as f32],
px as f32,
radius as f32,
[m.width as u32, m.height as u32],
)
}
None => (
self.label_texture(1, 1, &[dr_pano::seam::NONE]),
[0.0; 2],
1.0,
1.0,
[1, 1],
),
};
let seam_view = seam_tex.create_view(&Default::default());
let seam_on = u32::from(output.seams.is_some());
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,
clip_onset: dr_pipeline::CLIP_ONSET,
balance: [
output.balance[0].max(1e-3),
output.balance[1].max(1e-3),
output.balance[2].max(1e-3),
0.0,
],
seam_origin,
seam_size,
seam_px,
seam_radius,
frame_index: k as u32,
seam_on,
};
self.accumulate(&params, tile, &seam_view);
}
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()));
}
/// The seam map's labels as an `r8uint` texture.
fn label_texture(&self, width: u32, height: u32, labels: &[u8]) -> wgpu::Texture {
let size = wgpu::Extent3d {
width,
height,
depth_or_array_layers: 1,
};
let tex = self.ctx.device.create_texture(&wgpu::TextureDescriptor {
label: Some("merge-seams"),
size,
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: wgpu::TextureFormat::R8Uint,
usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
view_formats: &[],
});
self.ctx.queue.write_texture(
wgpu::TexelCopyTextureInfo {
texture: &tex,
mip_level: 0,
origin: wgpu::Origin3d::ZERO,
aspect: wgpu::TextureAspect::All,
},
labels,
wgpu::TexelCopyBufferLayout {
offset: 0,
bytes_per_row: Some(width),
rows_per_image: Some(height),
},
size,
);
tex
}
fn accumulate(&mut self, params: &WarpParams, tile: &wgpu::Texture, seams: &wgpu::TextureView) {
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(),
},
wgpu::BindGroupEntry {
binding: 3,
resource: wgpu::BindingResource::TextureView(seams),
},
],
});
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(&params),
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))
}