dtourolle
104e3a106f
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.
108 lines
4.2 KiB
Rust
108 lines
4.2 KiB
Rust
//! 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)
|
||
}
|
||
}
|