From 104e3a106f99c1a444d9a1fa948924dbcdab1cd2 Mon Sep 17 00:00:00 2001 From: Duncan Tourolle Date: Sat, 19 Sep 2026 20:33:51 +0200 Subject: [PATCH] Fill a panorama's border with MI-GAN: mirrored context, coarse to fine, a feathered seam MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit dr_pano::fill owns everything the model does not — which tiles, what context, how to blend — behind an Inpainter trait, and dr_pano::migan is that trait over the shipped generator on the inference engine. The known content is mirrored across the coverage edge into the hole and a 256-px ring, the nearest 48 px folded, so the model interpolates between real and mirrored sky rather than extrapolating into nothing. A coarse pass at a quarter decides the structure with the whole border in a few tiles; fine passes in 96-px bands from the edge outward texture it; the seam is blended over a feather inside the real edge. Every knob is a Params field, and an Observer hears each stage for whoever is looking at why a fill went wrong. --- core/dr-pano/Cargo.toml | 4 +- core/dr-pano/src/fill.rs | 793 ++++++++++++++++++++++++++++++++++++++ core/dr-pano/src/lib.rs | 11 + core/dr-pano/src/migan.rs | 107 +++++ docs/traceability.md | 6 +- 5 files changed, 916 insertions(+), 5 deletions(-) create mode 100644 core/dr-pano/src/fill.rs create mode 100644 core/dr-pano/src/migan.rs diff --git a/core/dr-pano/Cargo.toml b/core/dr-pano/Cargo.toml index f8088be..2499471 100644 --- a/core/dr-pano/Cargo.toml +++ b/core/dr-pano/Cargo.toml @@ -29,8 +29,8 @@ env_logger.workspace = true [features] default = ["xfeat", "embedded-model"] -# The XFeat detector (FR-MRG-8). Off, the crate has no model and no runtime, -# and `Detector` has no implementation — a build that only wants the geometry. +# The XFeat detector (FR-MRG-8) and the MI-GAN filler (FR-MRG-4). Off, the +# crate has no model and no runtime — a build that only wants the geometry. xfeat = ["dep:ort", "dep:dr-inference-engine", "dep:ndarray"] # Compile the weights into the binary, for the same reason `dr-segment` does: diff --git a/core/dr-pano/src/fill.rs b/core/dr-pano/src/fill.rs new file mode 100644 index 0000000..cbe1219 --- /dev/null +++ b/core/dr-pano/src/fill.rs @@ -0,0 +1,793 @@ +//! TRACES: FR-MRG-4 +//! Filling a composite's uncovered border, tile by tile, with an inpainter. +//! +//! A merged panorama has a ragged border where no frame reached. FR-MRG-4 +//! crops it by default; this fills it instead, when the photographer asks, +//! with pixels a model invents from the picture around them. Everything +//! here is the geometry of that — which tiles to run, what context to hand +//! the model, how to put its answers back — and none of it is the model: +//! that is the [`Inpainter`] trait, with MI-GAN behind it in `migan.rs` +//! and a fake in the tests. +//! +//! # Context across the edge +//! +//! An inpainting model is trained on holes *inside* pictures. A panorama's +//! border is a hole at the picture's *edge*: real content on one side, +//! nothing on the other, and a model given that invents a structure along +//! the open side — streaks of road in the sky, on the first try +//! (2026-09-19). So the known content is mirrored across the coverage +//! edge, column by column for the top and bottom bands and row by row for +//! the sides, into the hole and into a padding ring around the picture, +//! and the ring is presented as *known*. The model then interpolates +//! between real content and its mirror rather than extrapolating into +//! nothing. The ring is cut off at the end. +//! +//! # Structure from far away, texture from near +//! +//! One tiled pass at the working resolution was not enough: a 512-px tile +//! straddling the coverage edge sees a few hundred pixels of real content +//! on one side and invents the rest from that, two neighbouring tiles +//! invent differently, and the seams and the merge's own fringe leak into +//! the fill. [`fill_border`] therefore runs in two stages. A **coarse** +//! pass at a quarter of the size, where the whole border and hundreds of +//! pixels of real context sit inside a handful of tiles, decides the +//! structure — where the slope goes, where the sky stays sky. Then +//! **fine** passes regenerate the hole in bands from the real edge +//! outward: each band is the only unknown, with real content (or the band +//! before, freshly textured) on its near side and the coarse fill, +//! upsampled, on its far side — blurry, but the right structure — so the +//! model generates texture and a transition, never a large hole from +//! nothing. +//! +//! Tiles overlap by a third and are blended under a raised-cosine window, +//! so the seams between tiles do not show; the model's answer replaces +//! only the pixels that were unknown, and the picture itself is untouched. + +use crate::PanoError; + +/// A model that fills a square hole from its surroundings. +pub trait Inpainter { + /// The square tile it takes, in pixels. + fn tile(&self) -> usize; + + /// Fill one tile. `rgb` is `tile × tile × 3`, row-major, 0..1, with the + /// unknown pixels' values meaningless; `known` is `tile × tile`. The + /// result is `tile × tile × 3`, 0..1, of which only the unknown pixels + /// are read. + fn fill(&mut self, rgb: &[f32], known: &[bool]) -> Result, PanoError>; +} + +/// What a caller hears from [`fill_border`]: progress, for a page's bar, +/// and — for whoever is looking at why a fill went wrong — each stage's +/// picture as it lands. A plain `FnMut(usize, usize)` is an observer that +/// hears only the progress. +pub trait Observer { + /// `(done, total)` tiles, the total an estimate until the last band. + fn progress(&mut self, done: usize, total: usize); + /// A stage's result, `width × height × 3`: `coarse` (at the coarse + /// size), `band-N` after each fine band, `feathered` at the end. + fn stage(&mut self, _name: &str, _rgb: &[f32], _width: usize, _height: usize) {} +} + +impl Observer for F { + fn progress(&mut self, done: usize, total: usize) { + self(done, total) + } +} + +/// How far the picture is extended with mirrored content before tiling. +/// Half a tile: enough that a hole at the edge sits well inside a tile. +pub const RING: usize = 256; + +/// The fill's knobs, in pixels of the working image. The defaults are +/// what the fixture panorama looked best with on 2026-09-19; the merge +/// page exposes every one of them while the fill is experimental, so a +/// bad corner can be worked on from the picture rather than the code. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct Params { + /// The coarse pass's reduction: 1 skips it. + pub coarse: usize, + /// The fine passes' band width. + pub band: usize, + /// How deep into the picture the mirrored context reaches. A plain + /// reflection of a deep hole pulls in whatever is that far from the + /// edge — a ridge, a peak — and the model, told that is what lies + /// beyond, paints it upside down. Folding the reflection within this + /// band keeps the ring looking like the edge it continues (sky beside + /// sky, grass beside grass) and nothing further away. + pub mirror_depth: usize, + /// How far inside the real edge the fill also regenerates, the two + /// blended by distance. A hard cut between real pixels and invented + /// ones is a line whatever the fill's quality; blended over this many + /// pixels it is not. Zero is the hard cut. + pub feather: usize, + /// The step between tiles, at most the tile; two thirds of it usual. + pub stride: usize, +} + +impl Default for Params { + fn default() -> Self { + Params { + coarse: 4, + band: 96, + mirror_depth: 48, + feather: 24, + stride: 384, + } + } +} + +/// Fill the unknown pixels of `rgb` (`width × height × 3`, 0..1) in place: +/// the coarse pass, then the fine bands, then the seam feathered over +/// `feather` pixels inside the real edge. Returns the tiles run. +/// +/// `known` is `width × height`. `observer` hears the progress and, if it +/// cares, each stage. +pub fn fill_border( + rgb: &mut [f32], + width: usize, + height: usize, + known: &[bool], + model: &mut dyn Inpainter, + params: Params, + observer: &mut dyn Observer, +) -> Result { + let Params { + coarse: q, + band, + mirror_depth, + feather, + stride, + } = params; + let q = q.max(1); + let band = band.max(8); + let mirror_depth = mirror_depth.max(1); + if width == 0 || height == 0 || rgb.len() != width * height * 3 || known.len() != width * height + { + return Err(PanoError::Input("fill: buffer sizes disagree".into())); + } + if known.iter().all(|&k| k) { + return Ok(0); + } + // The fill regenerates a margin inside the real edge too, and the + // result is blended with the real pixels across it at the end. + let real = rgb.to_vec(); + let outer = known.to_vec(); + let mut inner = known.to_vec(); + erode(&mut inner, width, height, feather); + let known = &inner[..]; + let mut done = 0usize; + + // Coarse: a fraction of the size, unknown where any pixel of the cell was. + let (cw, ch) = ((width / q).max(1), (height / q).max(1)); + let mut coarse = vec![0.0f32; cw * ch * 3]; + let mut cknown = vec![true; cw * ch]; + for y in 0..ch { + for x in 0..cw { + let mut sum = [0.0f32; 3]; + let mut n = 0.0f32; + let mut all_known = true; + for dy in 0..q { + for dx in 0..q { + let (sx, sy) = ((x * q + dx).min(width - 1), (y * q + dy).min(height - 1)); + let i = sy * width + sx; + all_known &= known[i]; + for c in 0..3 { + sum[c] += rgb[i * 3 + c]; + } + n += 1.0; + } + } + for c in 0..3 { + coarse[(y * cw + x) * 3 + c] = sum[c] / n; + } + cknown[y * cw + x] = all_known; + } + } + let estimate = |tiles: usize| tiles * 4; + done += fill_once( + &mut coarse, + cw, + ch, + &cknown, + model, + stride, + mirror_depth, + |n, t| observer.progress(n, estimate(t)), + )?; + observer.stage("coarse", &coarse, cw, ch); + + // The hole starts as the coarse structure, upsampled. + for y in 0..height { + for x in 0..width { + let i = y * width + x; + if known[i] { + continue; + } + let fx = ((x as f32 + 0.5) / q as f32 - 0.5).clamp(0.0, (cw - 1) as f32); + let fy = ((y as f32 + 0.5) / q as f32 - 0.5).clamp(0.0, (ch - 1) as f32); + let (x0, y0) = (fx as usize, fy as usize); + let (x1, y1) = ((x0 + 1).min(cw - 1), (y0 + 1).min(ch - 1)); + let (tx, ty) = (fx - x0 as f32, fy - y0 as f32); + for c in 0..3 { + let at = |xx: usize, yy: usize| coarse[(yy * cw + xx) * 3 + c]; + rgb[i * 3 + c] = (at(x0, y0) * (1.0 - tx) + at(x1, y0) * tx) * (1.0 - ty) + + (at(x0, y1) * (1.0 - tx) + at(x1, y1) * tx) * ty; + } + } + } + + // Fine, in bands from the edge outward. + let dist = distance_to_known(known, width, height); + let mut band_known = vec![true; width * height]; + let mut b = 0usize; + loop { + let lo = (b * band).saturating_sub(band / 2) as f32; + let hi = ((b + 1) * band) as f32; + let mut any = false; + for i in 0..width * height { + let in_band = !known[i] && dist[i] > lo && dist[i] <= hi; + band_known[i] = !in_band; + any |= in_band; + } + if !any { + break; + } + let before = done; + done += fill_once( + rgb, + width, + height, + &band_known, + model, + stride, + mirror_depth, + |n, t| observer.progress(before + n, before + estimate(t)), + )?; + observer.stage(&format!("band-{b}"), rgb, width, height); + b += 1; + } + + // The seam: across the margin, real on the inside, invented on the + // outside, a smooth ramp between by distance from the true hole. + if feather > 0 { + let to_hole = + distance_to_known(&outer.iter().map(|k| !k).collect::>(), width, height); + for i in 0..width * height { + if !outer[i] || known[i] { + continue; + } + // In the margin: outer says known, inner says not. + let t = (to_hole[i] / feather as f32).clamp(0.0, 1.0); + let t = t * t * (3.0 - 2.0 * t); + for c in 0..3 { + rgb[i * 3 + c] = rgb[i * 3 + c] * (1.0 - t) + real[i * 3 + c] * t; + } + } + } + observer.stage("feathered", rgb, width, height); + observer.progress(done, done); + Ok(done) +} + +/// One tiled pass: every unknown pixel regenerated from the tiles that +/// touch it, the rest kept. Returns the tiles run. +#[allow(clippy::too_many_arguments)] +fn fill_once( + rgb: &mut [f32], + width: usize, + height: usize, + known: &[bool], + model: &mut dyn Inpainter, + stride: usize, + mirror_depth: usize, + mut progress: impl FnMut(usize, usize), +) -> Result { + let t = model.tile(); + if t == 0 || known.iter().all(|&k| k) { + return Ok(0); + } + + // The padded canvas with mirrored context, and the hole within it. + let ctx = MirroredContext::build(rgb, width, height, known, mirror_depth); + let (pw, ph) = (ctx.width, ctx.height); + + // Tiles that touch the hole, on a grid that reaches both far edges. + let starts = |n: usize| -> Vec { + if n <= t { + return vec![0]; + } + let mut v: Vec = (0..=n - t).step_by(stride.clamp(1, t)).collect(); + if *v.last().unwrap_or(&0) != n - t { + v.push(n - t); + } + v + }; + let ys = starts(ph); + let xs = starts(pw); + let mut tiles = Vec::new(); + for &y in &ys { + for &x in &xs { + if y + t > ph || x + t > pw { + continue; + } + let touches = + (y..y + t).any(|yy| ctx.hole[yy * pw + x..yy * pw + x + t].iter().any(|&h| h)); + if touches { + tiles.push((x, y)); + } + } + } + + // Raised-cosine window, so overlapping tiles blend. + let hann: Vec = (0..t) + .map(|i| { + let s = ((i as f32 + 1.0) / (t as f32 + 1.0) * std::f32::consts::PI).sin(); + s * s + 1e-3 + }) + .collect(); + + let mut acc = vec![0.0f32; pw * ph * 3]; + let mut wsum = vec![0.0f32; pw * ph]; + let mut tile_rgb = vec![0.0f32; t * t * 3]; + let mut tile_known = vec![false; t * t]; + let total = tiles.len(); + for (n, &(x, y)) in tiles.iter().enumerate() { + progress(n, total); + for r in 0..t { + let src = ((y + r) * pw + x) * 3; + tile_rgb[r * t * 3..(r + 1) * t * 3].copy_from_slice(&ctx.rgb[src..src + t * 3]); + let ks = (y + r) * pw + x; + for c in 0..t { + tile_known[r * t + c] = !ctx.hole[ks + c]; + } + } + let out = model.fill(&tile_rgb, &tile_known)?; + if out.len() != t * t * 3 { + return Err(PanoError::Model(format!( + "the inpainter returned {} values for a {t}×{t} tile", + out.len() + ))); + } + for r in 0..t { + for c in 0..t { + let w = hann[r] * hann[c]; + let p = (y + r) * pw + (x + c); + for ch in 0..3 { + acc[p * 3 + ch] += out[(r * t + c) * 3 + ch] * w; + } + wsum[p] += w; + } + } + } + progress(total, total); + + // Back into the picture: only the unknown pixels change. + for yy in 0..height { + for xx in 0..width { + let i = yy * width + xx; + if known[i] { + continue; + } + let p = (yy + RING) * pw + (xx + RING); + if wsum[p] > 0.0 { + for ch in 0..3 { + rgb[i * 3 + ch] = (acc[p * 3 + ch] / wsum[p]).clamp(0.0, 1.0); + } + } + } + } + Ok(total) +} + +/// Shrink `known` by `iterations` pixels on every side, in place. +/// +/// The merge's coverage edge carries a fringe — the last partly-covered +/// pixels of a frame, and whatever the renderer did at the boundary — and +/// a fill that stops exactly at the coverage bit leaves it as a dark line +/// along the seam. Eight pixels at a quarter of the composite's resolution +/// was what it took on the fixture. +pub fn erode(known: &mut [bool], width: usize, height: usize, iterations: usize) { + let mut next = known.to_vec(); + for _ in 0..iterations { + for y in 0..height { + for x in 0..width { + let i = y * width + x; + if !known[i] { + continue; + } + let edge = x == 0 + || y == 0 + || x + 1 == width + || y + 1 == height + || !known[i - 1] + || !known[i + 1] + || !known[i - width] + || !known[i + width]; + next[i] = !edge; + } + } + known.copy_from_slice(&next); + } +} + +/// Distance from each pixel to the nearest known one, by two chamfer +/// sweeps — within a few percent of Euclidean, and enough to cut bands. +fn distance_to_known(known: &[bool], width: usize, height: usize) -> Vec { + let inf = (width + height) as f32; + let mut d: Vec = known.iter().map(|&k| if k { 0.0 } else { inf }).collect(); + let (a, b) = (1.0f32, std::f32::consts::SQRT_2); + for y in 0..height { + for x in 0..width { + let i = y * width + x; + let mut v = d[i]; + if x > 0 { + v = v.min(d[i - 1] + a); + } + if y > 0 { + v = v.min(d[i - width] + a); + if x > 0 { + v = v.min(d[i - width - 1] + b); + } + if x + 1 < width { + v = v.min(d[i - width + 1] + b); + } + } + d[i] = v; + } + } + for y in (0..height).rev() { + for x in (0..width).rev() { + let i = y * width + x; + let mut v = d[i]; + if x + 1 < width { + v = v.min(d[i + 1] + a); + } + if y + 1 < height { + v = v.min(d[i + width] + a); + if x + 1 < width { + v = v.min(d[i + width + 1] + b); + } + if x > 0 { + v = v.min(d[i + width - 1] + b); + } + } + d[i] = v; + } + } + d +} + +/// Distance beyond the edge to distance inside it, folded within `depth` +/// ([`Params::mirror_depth`]): a triangle wave, so the band is read +/// forward and back rather than clamped to one row. +fn fold(d: usize, depth: usize) -> usize { + let period = 2 * depth; + let r = d % period; + if r <= depth { + r + } else { + period - r + } +} + +/// The picture on a canvas `RING` wider on every side, with the hole and +/// the ring filled by mirroring the known content across the coverage +/// edge — the nearest `depth` of it, folded — and the hole, the +/// original unknown and nothing else, marked. +struct MirroredContext { + width: usize, + height: usize, + rgb: Vec, + hole: Vec, +} + +impl MirroredContext { + fn build(rgb: &[f32], width: usize, height: usize, known: &[bool], depth: usize) -> Self { + let fold = |d: usize| fold(d, depth); + let (pw, ph) = (width + 2 * RING, height + 2 * RING); + let mut canvas = vec![0.0f32; pw * ph * 3]; + let mut kn = vec![false; pw * ph]; + let mut hole = vec![false; pw * ph]; + for y in 0..height { + for x in 0..width { + let i = y * width + x; + let p = (y + RING) * pw + (x + RING); + canvas[p * 3..p * 3 + 3].copy_from_slice(&rgb[i * 3..i * 3 + 3]); + kn[p] = known[i]; + hole[p] = !known[i]; + } + } + + // Per column: mirror across the first and last known row. + for x in 0..pw { + let first = (0..ph).find(|&y| kn[y * pw + x]); + let Some(first) = first else { continue }; + let last = (0..ph).rev().find(|&y| kn[y * pw + x]).unwrap_or(first); + for y in 0..first { + let m = (first + fold(first - y)).min(last); + let (d, s) = ((y * pw + x) * 3, (m * pw + x) * 3); + canvas.copy_within(s..s + 3, d); + } + for y in last + 1..ph { + let m = last.saturating_sub(fold(y - last)).max(first); + let (d, s) = ((y * pw + x) * 3, (m * pw + x) * 3); + canvas.copy_within(s..s + 3, d); + } + } + // Per row, for the sides, over what is there now. + for y in 0..ph { + let first = (0..pw).find(|&x| kn[y * pw + x]); + let Some(first) = first else { continue }; + let last = (0..pw).rev().find(|&x| kn[y * pw + x]).unwrap_or(first); + for x in 0..first { + let m = (first + fold(first - x)).min(last); + let (d, s) = ((y * pw + x) * 3, (y * pw + m) * 3); + canvas.copy_within(s..s + 3, d); + } + for x in last + 1..pw { + let m = last.saturating_sub(fold(x - last)).max(first); + let (d, s) = ((y * pw + x) * 3, (y * pw + m) * 3); + canvas.copy_within(s..s + 3, d); + } + } + MirroredContext { + width: pw, + height: ph, + rgb: canvas, + hole, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The tests' small pictures: a 48-px stride, a given feather. + fn test_params(feather: usize) -> Params { + Params { + stride: 48, + feather, + ..Params::default() + } + } + + /// Paints every unknown pixel a fixed grey and copies the known ones, + /// and remembers what it was shown. + struct Flat { + tile: usize, + seen: Vec<(Vec, Vec)>, + } + + impl Inpainter for Flat { + fn tile(&self) -> usize { + self.tile + } + fn fill(&mut self, rgb: &[f32], known: &[bool]) -> Result, PanoError> { + self.seen.push((rgb.to_vec(), known.to_vec())); + Ok(rgb + .chunks_exact(3) + .zip(known) + .flat_map(|(p, &k)| if k { [p[0], p[1], p[2]] } else { [0.5; 3] }) + .collect()) + } + } + + fn picture(w: usize, h: usize, border: usize) -> (Vec, Vec) { + let mut rgb = vec![0.0; w * h * 3]; + let mut known = vec![false; w * h]; + for y in 0..h { + for x in 0..w { + let i = y * w + x; + if y >= border && y < h - border { + known[i] = true; + rgb[i * 3] = x as f32 / w as f32; + rgb[i * 3 + 1] = y as f32 / h as f32; + rgb[i * 3 + 2] = 0.25; + } + } + } + (rgb, known) + } + + #[test] + fn unknown_pixels_take_the_model_and_known_ones_do_not_move() { + let (mut rgb, known) = picture(300, 200, 20); + let before = rgb.clone(); + let mut model = Flat { + tile: 64, + seen: Vec::new(), + }; + let tiles = fill_border( + &mut rgb, + 300, + 200, + &known, + &mut model, + test_params(0), + &mut |_, _| {}, + ) + .unwrap(); + assert!(tiles > 0); + for i in 0..300 * 200 { + if known[i] { + assert_eq!(&rgb[i * 3..i * 3 + 3], &before[i * 3..i * 3 + 3]); + } else { + for c in 0..3 { + assert!((rgb[i * 3 + c] - 0.5).abs() < 1e-4, "pixel {i}"); + } + } + } + } + + #[test] + fn the_model_is_shown_mirrored_context_not_black() { + let (mut rgb, known) = picture(300, 200, 20); + let mut model = Flat { + tile: 64, + seen: Vec::new(), + }; + fill_border( + &mut rgb, + 300, + 200, + &known, + &mut model, + test_params(0), + &mut |_, _| {}, + ) + .unwrap(); + for (tile_rgb, tile_known) in &model.seen { + let known_non_black = tile_rgb + .chunks_exact(3) + .zip(tile_known) + .filter(|(_, &k)| k) + .any(|(p, _)| p.iter().any(|v| *v > 0.0)); + assert!(known_non_black); + } + } + + #[test] + fn the_fine_passes_run_in_bands_after_the_coarse_one() { + // A 150-tall hole above and below a picture: the coarse pass sees + // it at a quarter; the fine passes need two bands of BAND pixels. + let (mut rgb, known) = picture(200, 500, 150); + let mut model = Flat { + tile: 64, + seen: Vec::new(), + }; + fill_border( + &mut rgb, + 200, + 500, + &known, + &mut model, + test_params(0), + &mut |_, _| {}, + ) + .unwrap(); + assert!(model.seen.len() > 4); + // Every unknown pixel was reached. + for i in 0..200 * 500 { + if !known[i] { + assert!((rgb[i * 3] - 0.5).abs() < 1e-4, "pixel {i}"); + } + } + } + + #[test] + fn the_seam_ramps_from_real_to_invented_across_the_feather() { + let (mut rgb, known) = picture(300, 200, 20); + let before = rgb.clone(); + let mut model = Flat { + tile: 64, + seen: Vec::new(), + }; + fill_border( + &mut rgb, + 300, + 200, + &known, + &mut model, + test_params(8), + &mut |_, _| {}, + ) + .unwrap(); + // Row 20 is the real edge; the margin runs to row 27. At the edge + // the value is the model's grey, eight rows in it is the picture's. + let at = |y: usize| rgb[(y * 300 + 150) * 3 + 2]; + assert!((at(20) - 0.5).abs() < 0.05, "{}", at(20)); + assert!((at(29) - before[(29 * 300 + 150) * 3 + 2]).abs() < 1e-4); + let (lo, hi) = (at(20).min(at(29)), at(20).max(at(29))); + assert!( + at(23) > lo + 0.02 && at(23) < hi - 0.02, + "{} between {lo} and {hi}", + at(23) + ); + // The hole itself is the model's. + assert!((at(5) - 0.5).abs() < 1e-4); + } + + #[test] + fn erosion_shrinks_the_known_region_from_every_edge() { + let (_, mut known) = picture(20, 20, 4); + erode(&mut known, 20, 20, 2); + assert!(known[8 * 20 + 10]); + assert!(!known[5 * 20 + 10]); + assert!(!known[8 * 20 + 1]); + } + + #[test] + fn distance_counts_pixels_from_the_known_region() { + let (_, known) = picture(20, 20, 4); + let d = distance_to_known(&known, 20, 20); + assert_eq!(d[4 * 20 + 10], 0.0); + assert!((d[3 * 20 + 10] - 1.0).abs() < 1e-6); + assert!((d[10] - 4.0).abs() < 1e-6); + } + + #[test] + fn a_fully_covered_picture_runs_nothing() { + let (mut rgb, known) = picture(100, 100, 0); + let mut model = Flat { + tile: 64, + seen: Vec::new(), + }; + assert_eq!( + fill_border( + &mut rgb, + 100, + 100, + &known, + &mut model, + test_params(0), + &mut |_, _| {} + ) + .unwrap(), + 0 + ); + } + + #[test] + fn the_context_mirrors_the_top_rows_upward() { + let (rgb, known) = picture(40, 30, 5); + let ctx = MirroredContext::build(&rgb, 40, 30, &known, 48); + let x = RING + 10; + let first = RING + 5; + for k in 1..=4 { + let above = ((first - k) * ctx.width + x) * 3; + let mirror = ((first + k) * ctx.width + x) * 3; + assert_eq!(&ctx.rgb[above..above + 3], &ctx.rgb[mirror..mirror + 3]); + } + assert!(ctx.hole[(RING + 2) * ctx.width + x]); + assert!(!ctx.hole[(RING - 2) * ctx.width + x]); + } + + #[test] + fn the_mirror_reaches_no_deeper_than_its_band() { + // A ridge 200 rows in must not appear in the ring: beyond the band + // the reflection folds back towards the edge rather than on into + // the picture. + let (mut rgb, known) = picture(40, 400, 5); + let ridge = 5 + 200; + for x in 0..40 { + rgb[(ridge * 40 + x) * 3..(ridge * 40 + x) * 3 + 3].copy_from_slice(&[0.9, 0.1, 0.1]); + } + let ctx = MirroredContext::build(&rgb, 40, 400, &known, 48); + let x = RING + 10; + for y in 0..RING + 5 { + let p = (y * ctx.width + x) * 3; + assert!( + ctx.rgb[p] < 0.5, + "row {y} of the ring shows the ridge ({:?})", + &ctx.rgb[p..p + 3] + ); + } + assert_eq!(fold(0, 48), 0); + assert_eq!(fold(48, 48), 48); + assert_eq!(fold(58, 48), 38); + assert_eq!(fold(96, 48), 0); + assert_eq!(fold(99, 48), 3); + } +} diff --git a/core/dr-pano/src/lib.rs b/core/dr-pano/src/lib.rs index 97d58f1..2ad8db9 100644 --- a/core/dr-pano/src/lib.rs +++ b/core/dr-pano/src/lib.rs @@ -35,10 +35,13 @@ pub mod align; pub mod bundle; pub mod features; +pub mod fill; pub mod homography; pub mod image; pub mod linalg; pub mod matching; +#[cfg(feature = "xfeat")] +pub mod migan; pub mod projection; #[cfg(feature = "xfeat")] pub mod xfeat; @@ -46,6 +49,7 @@ pub mod xfeat; pub use align::{align, AlignOptions, Alignment, Link, Unaligned}; pub use bundle::Cameras; pub use features::{Features, Keypoint}; +pub use fill::{fill_border, Inpainter, Observer, Params as FillParams}; pub use image::Gray; pub use projection::Projection; @@ -64,6 +68,13 @@ pub enum PanoError { Inference(#[source] ort::Error), } +#[cfg(feature = "xfeat")] +impl From for PanoError { + fn from(e: ort::Error) -> Self { + PanoError::Inference(e) + } +} + #[cfg(feature = "xfeat")] impl From for PanoError { fn from(e: dr_inference_engine::Error) -> Self { diff --git a/core/dr-pano/src/migan.rs b/core/dr-pano/src/migan.rs new file mode 100644 index 0000000..ca1fae8 --- /dev/null +++ b/core/dr-pano/src/migan.rs @@ -0,0 +1,107 @@ +//! TRACES: FR-MRG-4 +//! MI-GAN, the border filler, under the inference engine. +//! +//! Sargsyan et al., ICCV 2023 (Picsart AI Research): inpainting built for +//! phones — about six million parameters of plain convolutions, no FFT and +//! no attention, so it quantises and runs on a DSP. MIT, code and weights +//! (`models/LICENCE.md`). The bare 512 generator is what ships, exported +//! at a fixed shape by `tools/export-migan.sh`; its six operator types load +//! on every rung, and what they cost is the whole story of whether a fill +//! is interactive: 7.4 s a tile under tract, 0.4 s under ONNX Runtime's +//! CPU pool, 23 ms in fp16 and 13 ms in int8 on a laptop's TensorRT +//! (2026-09-19, docs/panorama.md §12). +//! +//! The model's contract, from the reference `export_inference_model.py`: +//! input `1×4×512×512` float — channel 0 is `mask − 0.5` with 1 where the +//! picture is known, channels 1–3 the RGB in −1..1 with the unknown pixels +//! zeroed; output `1×3×512×512` in −1..1, of which the caller keeps the +//! unknown pixels. That is [`crate::fill::Inpainter`], and the rest — +//! which tiles, what context, how to blend — is `fill.rs`. + +use crate::fill::Inpainter; +use crate::PanoError; + +/// The tile the shipped export takes. +pub const TILE: usize = 512; + +pub struct MiGan { + model: dr_inference_engine::Model, +} + +impl MiGan { + /// From the model file, in whichever form the engine's rung wants + /// (`resolve_model` picks an int8 sibling for the Hexagon). + pub fn from_path(path: &std::path::Path) -> Result { + use dr_inference_engine::{resolve_model, Role}; + let (path, form) = resolve_model(Role::Inpainter, path); + let bytes = std::fs::read(&path).map_err(PanoError::ModelRead)?; + Self::from_bytes(&bytes, form) + } + + pub fn from_bytes(bytes: &[u8], form: dr_inference_engine::Form) -> Result { + use dr_inference_engine::Role; + Ok(MiGan { + model: dr_inference_engine::open(Role::Inpainter, form, bytes)?, + }) + } + + /// Where the fill runs, for a status line. + pub fn rung(&self) -> Result { + Ok(self.model.acquire()?.rung()) + } +} + +impl Inpainter for MiGan { + fn tile(&self) -> usize { + TILE + } + + fn fill(&mut self, rgb: &[f32], known: &[bool]) -> Result, PanoError> { + let n = TILE * TILE; + if rgb.len() != n * 3 || known.len() != n { + return Err(PanoError::Input(format!( + "MI-GAN takes a {TILE}×{TILE} tile; given {} values and {} mask entries", + rgb.len(), + known.len() + ))); + } + // NCHW: the mask plane, then the three masked colour planes. + let mut input = vec![0.0f32; 4 * n]; + for i in 0..n { + let m = if known[i] { 1.0 } else { 0.0 }; + input[i] = m - 0.5; + for c in 0..3 { + input[(c + 1) * n + i] = (rgb[i * 3 + c] * 2.0 - 1.0) * m; + } + } + let tensor = ort::value::Tensor::from_array( + ndarray::Array::from_shape_vec(ndarray::IxDyn(&[1, 4, TILE, TILE]), input) + .expect("shape matches by construction"), + )?; + let started = std::time::Instant::now(); + let acquired = self.model.acquire()?; + let acquired_at = started.elapsed(); + let mut session = acquired.lock(); + let outputs = session.run(ort::inputs![tensor])?; + log::trace!( + "migan: tile on {} — acquire {:.1} ms, run {:.1} ms", + acquired.rung().label(), + acquired_at.as_secs_f64() * 1e3, + (started.elapsed() - acquired_at).as_secs_f64() * 1e3 + ); + let (shape, data) = outputs[0].try_extract_tensor::()?; + let dims: Vec = shape.iter().copied().collect(); + if dims != [1, 3, TILE as i64, TILE as i64] { + return Err(PanoError::Model(format!( + "MI-GAN output is {dims:?}, expected [1, 3, {TILE}, {TILE}]" + ))); + } + let mut out = vec![0.0f32; n * 3]; + for i in 0..n { + for c in 0..3 { + out[i * 3 + c] = (data[c * n + i] * 0.5 + 0.5).clamp(0.0, 1.0); + } + } + Ok(out) + } +} diff --git a/docs/traceability.md b/docs/traceability.md index 0f90992..000e49f 100644 --- a/docs/traceability.md +++ b/docs/traceability.md @@ -9,8 +9,8 @@ Denominators are parsed from [`requirements.md`](requirements.md) at run time, n | Metric | Value | |---|---| -| Source files scanned | 387 | -| TRACES tags found | 1558 | +| Source files scanned | 389 | +| TRACES tags found | 1560 | | Requirements defined | 184 | | Requirements deferred (post-v1) | 24 | | Requirements covered | 154 | @@ -105,7 +105,7 @@ A tag naming an ID `requirements.md` does not define — what renumbering produc | FR-MRG-11 | [`core/dr-gpu/src/merge.rs:1`](../core/dr-gpu/src/merge.rs#L1), [`core/dr-gpu/src/shaders/merge.wgsl:1`](../core/dr-gpu/src/shaders/merge.wgsl#L1) | | FR-MRG-2 | [`core/dr-gpu/src/adjust.rs:104`](../core/dr-gpu/src/adjust.rs#L104), [`core/dr-gpu/src/adjust.rs:1157`](../core/dr-gpu/src/adjust.rs#L1157), [`core/dr-gpu/src/adjust.rs:1276`](../core/dr-gpu/src/adjust.rs#L1276), [`core/dr-gpu/src/adjust.rs:316`](../core/dr-gpu/src/adjust.rs#L316), [`core/dr-pipeline/src/graph.rs:852`](../core/dr-pipeline/src/graph.rs#L852), [`core/dr-pipeline/src/operation.rs:2124`](../core/dr-pipeline/src/operation.rs#L2124), [`core/dr-pipeline/src/operation.rs:427`](../core/dr-pipeline/src/operation.rs#L427), [`core/dr-pipeline/src/operation.rs:595`](../core/dr-pipeline/src/operation.rs#L595), [`core/dr-pipeline/src/operation.rs:860`](../core/dr-pipeline/src/operation.rs#L860), [`ui/dr-ui/src/merge.rs:1`](../ui/dr-ui/src/merge.rs#L1) | | FR-MRG-3 | [`core/dr-decode/examples/linear_dng.rs:1`](../core/dr-decode/examples/linear_dng.rs#L1), [`core/dr-decode/src/lib.rs:137`](../core/dr-decode/src/lib.rs#L137), [`core/dr-decode/src/lib.rs:144`](../core/dr-decode/src/lib.rs#L144), [`core/dr-decode/src/lib.rs:578`](../core/dr-decode/src/lib.rs#L578), [`core/dr-decode/src/profile.rs:395`](../core/dr-decode/src/profile.rs#L395), [`core/dr-export/src/dng.rs:1`](../core/dr-export/src/dng.rs#L1), [`core/dr-gpu/src/demosaic.rs:254`](../core/dr-gpu/src/demosaic.rs#L254), [`ui/dr-ui/src/merge.rs:1`](../ui/dr-ui/src/merge.rs#L1) | -| FR-MRG-4 | [`core/dr-export/src/inscribed.rs:1`](../core/dr-export/src/inscribed.rs#L1), [`core/dr-pano/src/projection.rs:1`](../core/dr-pano/src/projection.rs#L1), [`ui/dr-ui/src/merge_ui.rs:1`](../ui/dr-ui/src/merge_ui.rs#L1), [`ui/dr-ui/ui/merge.slint:1`](../ui/dr-ui/ui/merge.slint#L1) | +| FR-MRG-4 | [`core/dr-export/src/inscribed.rs:1`](../core/dr-export/src/inscribed.rs#L1), [`core/dr-pano/src/fill.rs:1`](../core/dr-pano/src/fill.rs#L1), [`core/dr-pano/src/migan.rs:1`](../core/dr-pano/src/migan.rs#L1), [`core/dr-pano/src/projection.rs:1`](../core/dr-pano/src/projection.rs#L1), [`ui/dr-ui/src/merge_ui.rs:1`](../ui/dr-ui/src/merge_ui.rs#L1), [`ui/dr-ui/ui/merge.slint:1`](../ui/dr-ui/ui/merge.slint#L1) | | FR-MRG-5 | [`core/dr-pano/src/align.rs:1`](../core/dr-pano/src/align.rs#L1), [`ui/dr-ui/src/merge.rs:1`](../ui/dr-ui/src/merge.rs#L1), [`ui/dr-ui/src/merge_ui.rs:1`](../ui/dr-ui/src/merge_ui.rs#L1), [`ui/dr-ui/ui/merge.slint:1`](../ui/dr-ui/ui/merge.slint#L1) | | FR-MRG-6 | [`core/dr-pipeline/src/sidecar.rs:1033`](../core/dr-pipeline/src/sidecar.rs#L1033), [`core/dr-pipeline/src/sidecar.rs:113`](../core/dr-pipeline/src/sidecar.rs#L113), [`core/dr-pipeline/src/sidecar.rs:123`](../core/dr-pipeline/src/sidecar.rs#L123), [`core/dr-pipeline/src/sidecar.rs:2241`](../core/dr-pipeline/src/sidecar.rs#L2241), [`core/dr-pipeline/src/sidecar.rs:841`](../core/dr-pipeline/src/sidecar.rs#L841), [`ui/dr-ui/src/merge.rs:509`](../ui/dr-ui/src/merge.rs#L509) | | FR-MRG-7 | [`ui/dr-ui/src/merge.rs:1`](../ui/dr-ui/src/merge.rs#L1), [`ui/dr-ui/src/merge_ui.rs:1`](../ui/dr-ui/src/merge_ui.rs#L1), [`ui/dr-ui/ui/app.slint:828`](../ui/dr-ui/ui/app.slint#L828), [`ui/dr-ui/ui/merge.slint:1`](../ui/dr-ui/ui/merge.slint#L1) |