Fill a panorama's border with MI-GAN: mirrored context, coarse to fine, a feathered seam
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.
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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<Vec<f32>, 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<F: FnMut(usize, usize)> 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<usize, PanoError> {
|
||||
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::<Vec<_>>(), 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<usize, PanoError> {
|
||||
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<usize> {
|
||||
if n <= t {
|
||||
return vec![0];
|
||||
}
|
||||
let mut v: Vec<usize> = (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<f32> = (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<f32> {
|
||||
let inf = (width + height) as f32;
|
||||
let mut d: Vec<f32> = 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<f32>,
|
||||
hole: Vec<bool>,
|
||||
}
|
||||
|
||||
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<f32>, Vec<bool>)>,
|
||||
}
|
||||
|
||||
impl Inpainter for Flat {
|
||||
fn tile(&self) -> usize {
|
||||
self.tile
|
||||
}
|
||||
fn fill(&mut self, rgb: &[f32], known: &[bool]) -> Result<Vec<f32>, 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<f32>, Vec<bool>) {
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -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<ort::Error> for PanoError {
|
||||
fn from(e: ort::Error) -> Self {
|
||||
PanoError::Inference(e)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "xfeat")]
|
||||
impl From<dr_inference_engine::Error> for PanoError {
|
||||
fn from(e: dr_inference_engine::Error) -> Self {
|
||||
|
||||
@@ -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<Self, PanoError> {
|
||||
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<Self, PanoError> {
|
||||
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<dr_inference_engine::Rung, PanoError> {
|
||||
Ok(self.model.acquire()?.rung())
|
||||
}
|
||||
}
|
||||
|
||||
impl Inpainter for MiGan {
|
||||
fn tile(&self) -> usize {
|
||||
TILE
|
||||
}
|
||||
|
||||
fn fill(&mut self, rgb: &[f32], known: &[bool]) -> Result<Vec<f32>, 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::<f32>()?;
|
||||
let dims: Vec<i64> = 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)
|
||||
}
|
||||
}
|
||||
@@ -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) |
|
||||
|
||||
Reference in New Issue
Block a user