Files
DarkRoom/core/dr-pano/src/migan.rs
T
dtourolle 84fade99ec Put the developer docs under docs/dev and index the folder for users first
docs/ had 26 developer documents flat beside the manual, and the two
audiences are very differently sized: most readers want the manual and
the gesture reference, a few want the register, the designs and the
measurements. The manual and gestures.md stay at the top; everything for
someone changing the code moves to docs/dev/, and the two documents that
name their own successors — the v0.1 milestone and the UI-refinement plan
— go to docs/dev/archive/ rather than being deleted, since both are still
cited. docs/README.md is the index, users first.

Every reference follows: code comments, Cargo manifests, the workflows,
the pre-commit hook, the bench and traceability tools (which locate the
repo root by docs/dev/requirements.md now), packaging, the Docker READMEs,
CLAUDE.md, CONTRIBUTING.md and the README. The matrix links one level
deeper and is regenerated. Links out of the moved documents into the tree
gain a level; a link checker over every Markdown file finds none broken.
2026-09-20 21:16:03 +02:00

108 lines
4.2 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-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/dev/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)
}
}