dtourolle
5a8c3e4c40
The engine knew f32 and int8, and gave the Hexagon int8 for every role it served. Measured on the tablet itself (inference.md §1.5), int8 lost 5% of the detector's faces at 40-80 px, moved the landmarks 1.5 px, emptied the segmenter's scores and cost the denoiser 5-9 dB; fp16 the HTP refuses outright. `Form` gains A16W8 and A16W16, and `Rung::form` now names one per role: detectors and landmarks A16W8, the segmenter, scene model, border filler and denoiser A16W16, XFeat int8. The embedder and the eye classifiers stay on the CPU. Each loader resolves its `<stem>.<form>.onnx` sibling; the segmenter and XFeat, compiled into the binary, embed their quantised forms on Android only and pick through `choose_embedded`. The probe, the compile step and the cache fingerprint follow the form instead of assuming int8. Detectors on the new form write `scrfd_*_a16+w600k_mbf`, and `model_ids` answers for all three spellings. On the tablet (ORT 1.29 + QNN 2.42), each shipped file against f32 on the same inputs, and against the CPU's f32 time: SCRFD 500m/2.5g/10g A16W8 100% of faces in every band 4.2/5.1/9.0 ms vs 17/56/198 landmarks A16W8 0.25 px in the 192 crop 0.5 ms vs 2.8 YOLO26n-seg A16W16 98.2% found, mask IoU 0.994 12.9 ms vs 90 scene model A16W16 98.9% of cells agree 15 ms vs 151 MI-GAN A16W16 41 dB from f32 in the fill 87 ms vs 488 XFeat int8 pano alignment 0.45 px (f32's own spread 0.41) 6.5 ms vs 58 denoiser A16W16 0.00 dB at every ISO 95 ms vs 1510 a tile Face numbers are over public COCO val2017 photographs, not a library. The APK carries the siblings (BUNDLED 15 -> 19; the old int8 detectors removed), about 43 MB more. The Windows installer and its CI count skip them; the Arch and Flatpak packages list their files and never had them. The ladder example takes a role per model, which is how the per-role forms above were seen landing on the NPU from the real probe.
883 lines
33 KiB
Rust
883 lines
33 KiB
Rust
//! Semantic segmentation — arm B (S15, docs/dev/segmentation.md §4).
|
||
//!
|
||
//! Runs a YOLO instance-segmentation graph over a proxy-resolution image and
|
||
//! returns the instances it found: a class, a score, a box, and a soft mask
|
||
//! each. [`crate::prior`] is what turns those into a merge prior over the
|
||
//! watershed hierarchy; nothing here knows about regions.
|
||
//!
|
||
//! # This is instance segmentation, not semantic segmentation
|
||
//!
|
||
//! §4 of the spec assumed a *semantic* model — a full partition of the image
|
||
//! into 150 ADE20K classes, sky and vegetation among them. The model that
|
||
//! actually exists is COCO-trained and *instance*-based, and the difference is
|
||
//! not cosmetic:
|
||
//!
|
||
//! - **It does not partition the image.** It finds objects. Most pixels in a
|
||
//! landscape belong to no instance at all, and that is not a failure — there
|
||
//! is no COCO class for "hillside".
|
||
//! - **It separates two people**, where a semantic model would hand back one
|
||
//! "person" area covering both. For selecting a subject this is the better
|
||
//! behaviour, and it is worth being glad of rather than working around.
|
||
//!
|
||
//! So arm B here contributes *subjects*, and the watershed contributes
|
||
//! everything else. See `models/LICENCE.md` at the repository root for why no ADE20K variant is
|
||
//! shipped instead.
|
||
//!
|
||
//! # Cost, and where it may run
|
||
//!
|
||
//! ~470 ms for one 640×640 inference on the reference desktop's CPU, pure Rust
|
||
//! via tract. That is a **once-per-image background precompute** and nothing
|
||
//! else: it must never sit on the frame path (ARCH §6.1), and the interactive
|
||
//! operations it enables — click a subject, grow a selection — read its cached
|
||
//! output rather than re-running it.
|
||
|
||
use std::sync::Arc;
|
||
|
||
use ndarray::{Array4, ArrayView2, ArrayView3};
|
||
|
||
use crate::SegmentError;
|
||
|
||
/// The graph's fixed input edge, in pixels.
|
||
///
|
||
/// **Fixed, not configurable.** tract cannot parse the dynamic-shape export of
|
||
/// this model — it fails shape inference on the neck's `Concat` — so the graph
|
||
/// ships with its input baked to one square size. Everything else in this
|
||
/// module, letterboxing and tiling alike, exists to fit arbitrary images
|
||
/// through that fixed window.
|
||
pub const INPUT_EDGE: usize = 640;
|
||
|
||
/// Detections per forward pass, from the graph's output shape `[1, 300, 38]`.
|
||
const MAX_DETECTIONS: usize = 300;
|
||
/// Mask prototypes, from `[1, 32, 160, 160]`.
|
||
const PROTOTYPES: usize = 32;
|
||
/// `4` box + `1` score + `1` class + `PROTOTYPES` coefficients.
|
||
const DETECTION_STRIDE: usize = 6 + PROTOTYPES;
|
||
/// Prototype masks come out at a quarter of the input edge.
|
||
const PROTO_STRIDE: usize = 4;
|
||
|
||
/// How the image is presented to a fixed-shape graph.
|
||
#[derive(Debug, Clone, Copy, PartialEq, Default)]
|
||
pub enum Tiling {
|
||
/// One inference over the whole frame, letterboxed into the square input.
|
||
///
|
||
/// The default, and the right default: a photographic subject is usually
|
||
/// *large* in frame, which is the case whole-image inference handles best
|
||
/// and the case tiling helps least.
|
||
#[default]
|
||
Whole,
|
||
/// Cover the frame with overlapping fixed-size windows.
|
||
///
|
||
/// Buys resolution — a subject 200 px across in a 1600 px proxy reaches
|
||
/// the model at 200 px rather than at 80 — and costs one inference per
|
||
/// tile. Worth it for a small subject in a large frame (a bird against
|
||
/// sky, a figure in a landscape) and wasteful otherwise.
|
||
///
|
||
/// `overlap` is the fraction of a tile shared with its neighbour, which
|
||
/// has to exceed zero or a subject sitting on a seam is cut in half by
|
||
/// both tiles and recognised by neither.
|
||
Grid { overlap: f32 },
|
||
}
|
||
|
||
/// How the semantic pass is tuned.
|
||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||
pub struct SemanticOptions {
|
||
/// Minimum detection score to keep.
|
||
///
|
||
/// Deliberately low. A false positive costs a spurious entry in a list the
|
||
/// user is choosing from; a false negative costs a subject that cannot be
|
||
/// selected at all, which is the worse failure for a selection tool.
|
||
pub confidence: f32,
|
||
/// Mask probability above which a pixel is inside the instance.
|
||
pub mask_threshold: f32,
|
||
pub tiling: Tiling,
|
||
/// Mask IoU above which two detections from different tiles are judged to
|
||
/// be the same object. Unused when [`Tiling::Whole`].
|
||
pub merge_iou: f32,
|
||
}
|
||
|
||
impl Default for SemanticOptions {
|
||
fn default() -> Self {
|
||
Self {
|
||
confidence: 0.25,
|
||
mask_threshold: 0.5,
|
||
tiling: Tiling::Whole,
|
||
merge_iou: 0.55,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// One detected object.
|
||
#[derive(Debug, Clone, PartialEq)]
|
||
pub struct Instance {
|
||
pub class_id: u16,
|
||
pub class_name: Arc<str>,
|
||
pub score: f32,
|
||
/// Bounding box in **source image** pixels: `(x0, y0, x1, y1)`.
|
||
pub bbox: (f32, f32, f32, f32),
|
||
/// Per-pixel coverage over the whole source image, row-major, `0.0..=1.0`.
|
||
///
|
||
/// Soft rather than binary because arm C weights merges by it, and a hard
|
||
/// threshold there would throw away exactly the confidence information
|
||
/// that makes a prior a prior rather than a decision.
|
||
pub mask: Vec<f32>,
|
||
pub width: usize,
|
||
pub height: usize,
|
||
}
|
||
|
||
impl Instance {
|
||
/// Fraction of this instance's mass inside a set of pixels.
|
||
pub fn coverage(&self, pixels: impl Iterator<Item = usize>) -> f32 {
|
||
let mut inside = 0.0;
|
||
let mut n = 0usize;
|
||
for p in pixels {
|
||
inside += self.mask.get(p).copied().unwrap_or(0.0);
|
||
n += 1;
|
||
}
|
||
if n == 0 {
|
||
0.0
|
||
} else {
|
||
inside / n as f32
|
||
}
|
||
}
|
||
|
||
/// Take in another view of the same object, from an overlapping tile.
|
||
///
|
||
/// Pointwise maximum over the coverage, union of the boxes, higher of the
|
||
/// scores — see `merge_into` for why each is right for its own quantity.
|
||
/// Both instances describe the whole frame in the same coordinates, so no
|
||
/// resampling is involved and the masks are already aligned pixel for
|
||
/// pixel.
|
||
fn absorb(&mut self, other: Self) {
|
||
debug_assert_eq!(
|
||
(self.width, self.height),
|
||
(other.width, other.height),
|
||
"instances from one detection run share the frame they are defined over"
|
||
);
|
||
for (mine, theirs) in self.mask.iter_mut().zip(&other.mask) {
|
||
if *theirs > *mine {
|
||
*mine = *theirs;
|
||
}
|
||
}
|
||
self.bbox = (
|
||
self.bbox.0.min(other.bbox.0),
|
||
self.bbox.1.min(other.bbox.1),
|
||
self.bbox.2.max(other.bbox.2),
|
||
self.bbox.3.max(other.bbox.3),
|
||
);
|
||
if other.score > self.score {
|
||
self.score = other.score;
|
||
// The name travels with the score: they are one judgement, and a
|
||
// mask labelled by the less confident of two detections would be
|
||
// labelled by the one we just decided to trust less.
|
||
self.class_name = other.class_name;
|
||
}
|
||
}
|
||
|
||
fn iou(&self, other: &Self, threshold: f32) -> f32 {
|
||
let mut inter = 0usize;
|
||
let mut union = 0usize;
|
||
for (a, b) in self.mask.iter().zip(&other.mask) {
|
||
let (a, b) = (*a >= threshold, *b >= threshold);
|
||
inter += usize::from(a && b);
|
||
union += usize::from(a || b);
|
||
}
|
||
if union == 0 {
|
||
0.0
|
||
} else {
|
||
inter as f32 / union as f32
|
||
}
|
||
}
|
||
}
|
||
|
||
/// A loaded segmentation model.
|
||
///
|
||
/// Holds an `ort` session, so it is neither `Clone` nor cheap to build —
|
||
/// construct once and keep it. Loading is ~50 ms.
|
||
pub struct SemanticModel {
|
||
session: dr_inference_engine::Model,
|
||
classes: Vec<Arc<str>>,
|
||
}
|
||
|
||
/// The weights, from the repository-root `models/segment/` (AGPL — see `models/LICENCE.md`).
|
||
///
|
||
/// Embedded rather than read from a path because Android hands the app no
|
||
/// filesystem location to read from (ARCH §6.9) — the same reasoning that has
|
||
/// the Lensfun database shipping inside its crate.
|
||
#[cfg(feature = "embedded-model")]
|
||
const EMBEDDED_MODEL: &[u8] = include_bytes!("../../../models/segment/yolo26n-seg.onnx");
|
||
#[cfg(feature = "embedded-model")]
|
||
const EMBEDDED_CLASSES: &str = include_str!("../../../models/segment/yolo26n-seg.classes.json");
|
||
|
||
/// The Hexagon's form (docs/dev/inference.md §1.5): 16-bit activations and
|
||
/// weights, the rows' tail left in float. Only Android has a Hexagon, so only
|
||
/// Android carries it.
|
||
#[cfg(all(feature = "embedded-model", target_os = "android"))]
|
||
const EMBEDDED_A16W16: &[u8] = include_bytes!("../../../models/segment/yolo26n-seg.a16w16.onnx");
|
||
|
||
/// Every form of the model that ships with this crate, for whoever compiles
|
||
/// engines ahead of the first request (docs/dev/inference.md §6).
|
||
#[cfg(feature = "embedded-model")]
|
||
pub fn embedded_models() -> Vec<(dr_inference_engine::Form, &'static [u8])> {
|
||
#[allow(unused_mut)]
|
||
let mut forms = vec![(dr_inference_engine::Form::F32, EMBEDDED_MODEL)];
|
||
#[cfg(target_os = "android")]
|
||
forms.push((dr_inference_engine::Form::A16W16, EMBEDDED_A16W16));
|
||
forms
|
||
}
|
||
|
||
impl SemanticModel {
|
||
/// Load the model that ships with this crate, in the form the device's
|
||
/// backend runs.
|
||
#[cfg(feature = "embedded-model")]
|
||
pub fn embedded() -> Result<Self, SegmentError> {
|
||
let forms = embedded_models();
|
||
let (bytes, form) =
|
||
dr_inference_engine::choose_embedded(dr_inference_engine::Role::Segmenter, &forms);
|
||
Self::from_bytes_in(bytes, form, parse_classes(EMBEDDED_CLASSES))
|
||
}
|
||
|
||
/// Load a model from an ONNX file, with `classes` supplying its vocabulary.
|
||
///
|
||
/// The vocabulary is a parameter rather than a constant so that swapping in
|
||
/// a model with different classes — an ADE20K stuff model, say — is a data
|
||
/// change rather than a code change.
|
||
pub fn from_path(
|
||
path: impl AsRef<std::path::Path>,
|
||
classes: Vec<Arc<str>>,
|
||
) -> Result<Self, SegmentError> {
|
||
let bytes = std::fs::read(path).map_err(SegmentError::ModelRead)?;
|
||
Self::from_bytes(&bytes, classes)
|
||
}
|
||
|
||
pub fn from_bytes(bytes: &[u8], classes: Vec<Arc<str>>) -> Result<Self, SegmentError> {
|
||
Self::from_bytes_in(bytes, dr_inference_engine::Form::F32, classes)
|
||
}
|
||
|
||
/// `bytes` in a stated numeric form. The quantised one keeps the same
|
||
/// outputs (the rows' tail stays float), so decoding does not change; the
|
||
/// masks it draws were measured against f32's (inference.md §1.5).
|
||
pub fn from_bytes_in(
|
||
bytes: &[u8],
|
||
form: dr_inference_engine::Form,
|
||
classes: Vec<Arc<str>>,
|
||
) -> Result<Self, SegmentError> {
|
||
let session = dr_inference_engine::open(dr_inference_engine::Role::Segmenter, form, bytes)?;
|
||
|
||
Ok(Self { session, classes })
|
||
}
|
||
|
||
pub fn classes(&self) -> &[Arc<str>] {
|
||
&self.classes
|
||
}
|
||
|
||
/// Find the objects in an image.
|
||
///
|
||
/// `rgb` is tightly packed `f32` RGB in `0.0..=1.0`, row-major, three
|
||
/// components per pixel — the same linear-ish proxy the watershed reads,
|
||
/// so both arms describe the same picture.
|
||
pub fn detect(
|
||
&mut self,
|
||
rgb: &[f32],
|
||
width: usize,
|
||
height: usize,
|
||
options: &SemanticOptions,
|
||
) -> Result<Vec<Instance>, SegmentError> {
|
||
if width == 0 || height == 0 {
|
||
return Ok(Vec::new());
|
||
}
|
||
if rgb.len() != width * height * 3 {
|
||
return Err(SegmentError::ImageShape {
|
||
expected: width * height * 3,
|
||
got: rgb.len(),
|
||
});
|
||
}
|
||
|
||
let windows = self.windows(width, height, options.tiling);
|
||
let mut found: Vec<Instance> = Vec::new();
|
||
|
||
for window in &windows {
|
||
let batch = self.run_window(rgb, width, height, window, options)?;
|
||
merge_into(&mut found, batch, options);
|
||
}
|
||
|
||
// Strongest first: this list is offered to a person as "what did you
|
||
// mean", and the most confident guess belongs at the top.
|
||
found.sort_by(|a, b| b.score.total_cmp(&a.score));
|
||
Ok(found)
|
||
}
|
||
|
||
/// The source-space rectangles each inference covers.
|
||
fn windows(&self, width: usize, height: usize, tiling: Tiling) -> Vec<Window> {
|
||
match tiling {
|
||
Tiling::Whole => vec![Window {
|
||
x: 0.0,
|
||
y: 0.0,
|
||
w: width as f32,
|
||
h: height as f32,
|
||
}],
|
||
Tiling::Grid { overlap } => {
|
||
// A tile covers a square of source pixels whose edge is the
|
||
// shorter image dimension clamped to something the model can
|
||
// still see detail in. Below that the tiling is pointless —
|
||
// the window is already smaller than the input.
|
||
let edge = (width.min(height) as f32).min(INPUT_EDGE as f32 * 1.5);
|
||
let overlap = overlap.clamp(0.0, 0.9);
|
||
let stride = (edge * (1.0 - overlap)).max(1.0);
|
||
|
||
let mut windows = Vec::new();
|
||
for gy in 0..steps(height as f32, edge, stride) {
|
||
for gx in 0..steps(width as f32, edge, stride) {
|
||
// Last row and column are pulled back inside the frame
|
||
// rather than padded, so no inference is spent on
|
||
// blank margin.
|
||
let x = (gx as f32 * stride).min((width as f32 - edge).max(0.0));
|
||
let y = (gy as f32 * stride).min((height as f32 - edge).max(0.0));
|
||
windows.push(Window {
|
||
x,
|
||
y,
|
||
w: edge.min(width as f32),
|
||
h: edge.min(height as f32),
|
||
});
|
||
}
|
||
}
|
||
windows
|
||
}
|
||
}
|
||
}
|
||
|
||
fn run_window(
|
||
&mut self,
|
||
rgb: &[f32],
|
||
width: usize,
|
||
height: usize,
|
||
window: &Window,
|
||
options: &SemanticOptions,
|
||
) -> Result<Vec<Instance>, SegmentError> {
|
||
let letterbox = Letterbox::fit(window.w, window.h);
|
||
let input = letterbox.sample(rgb, width, height, window);
|
||
|
||
let acquired = self.session.acquire()?;
|
||
let mut session = acquired.lock();
|
||
let outputs = session
|
||
.run(ort::inputs![
|
||
ort::value::Tensor::from_array(input).map_err(SegmentError::Inference)?
|
||
])
|
||
.map_err(SegmentError::Inference)?;
|
||
|
||
let (det_shape, det) = outputs[0]
|
||
.try_extract_tensor::<f32>()
|
||
.map_err(SegmentError::Inference)?;
|
||
let (proto_shape, proto) = outputs[1]
|
||
.try_extract_tensor::<f32>()
|
||
.map_err(SegmentError::Inference)?;
|
||
|
||
// The decoder reads fixed column offsets out of each row, so a row of
|
||
// an unexpected width means a model this code cannot read — a
|
||
// different class count, a different prototype count, a detect-only
|
||
// export. Caught here as an error rather than downstream as garbage
|
||
// boxes, because garbage boxes look like a bad model rather than a
|
||
// wrong one.
|
||
if det_shape[2] as usize != DETECTION_STRIDE {
|
||
return Err(SegmentError::OutputShape("detections"));
|
||
}
|
||
|
||
let detections = ArrayView2::from_shape(
|
||
(det_shape[1] as usize, det_shape[2] as usize),
|
||
&det[..(det_shape[1] * det_shape[2]) as usize],
|
||
)
|
||
.map_err(|_| SegmentError::OutputShape("detections"))?;
|
||
|
||
let (pc, ph, pw) = (
|
||
proto_shape[1] as usize,
|
||
proto_shape[2] as usize,
|
||
proto_shape[3] as usize,
|
||
);
|
||
let protos = ArrayView3::from_shape((pc, ph, pw), &proto[..pc * ph * pw])
|
||
.map_err(|_| SegmentError::OutputShape("prototypes"))?;
|
||
|
||
// `&self.classes` rather than `self.decode(..)`: `outputs` holds a
|
||
// mutable borrow of `self.session` until it drops, and a method call
|
||
// would borrow all of `self`. Borrowing the two fields separately is
|
||
// what the borrow checker will actually allow here.
|
||
Ok(decode(
|
||
&self.classes,
|
||
detections,
|
||
protos,
|
||
&letterbox,
|
||
window,
|
||
width,
|
||
height,
|
||
options,
|
||
))
|
||
}
|
||
}
|
||
|
||
/// Turn one forward pass into instances in source-image space.
|
||
///
|
||
/// YOLO26 is **NMS-free**: the head emits a fixed 300 slots already suppressed
|
||
/// and score-ordered, so there is no non-maximum suppression to implement
|
||
/// here. Only the cross-*tile* duplicates need merging, and that is
|
||
/// [`merge_into`]'s job.
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn decode(
|
||
classes: &[Arc<str>],
|
||
detections: ArrayView2<f32>,
|
||
protos: ArrayView3<f32>,
|
||
letterbox: &Letterbox,
|
||
window: &Window,
|
||
width: usize,
|
||
height: usize,
|
||
options: &SemanticOptions,
|
||
) -> Vec<Instance> {
|
||
let (ph, pw) = (protos.shape()[1], protos.shape()[2]);
|
||
let mut out = Vec::new();
|
||
|
||
for d in 0..detections.shape()[0].min(MAX_DETECTIONS) {
|
||
let row = detections.row(d);
|
||
let score = row[4];
|
||
// Score-ordered, so the first miss ends the useful part of the batch
|
||
// and the remaining slots are padding.
|
||
if score < options.confidence {
|
||
break;
|
||
}
|
||
|
||
let class_id = row[5] as u16;
|
||
let Some(class_name) = classes.get(class_id as usize).cloned() else {
|
||
continue;
|
||
};
|
||
|
||
// Box is in letterboxed input space; undo the letterbox and the window
|
||
// offset to land in source pixels.
|
||
let bbox = letterbox.to_source(row[0], row[1], row[2], row[3], window);
|
||
|
||
let coeffs: Vec<f32> = row.iter().skip(6).take(PROTOTYPES).copied().collect();
|
||
let mask = assemble_mask(
|
||
&coeffs, protos, ph, pw, letterbox, window, &bbox, width, height, options,
|
||
);
|
||
|
||
out.push(Instance {
|
||
class_id,
|
||
class_name,
|
||
score,
|
||
bbox,
|
||
mask,
|
||
width,
|
||
height,
|
||
});
|
||
}
|
||
|
||
out
|
||
}
|
||
|
||
/// A source-space rectangle fed through one inference.
|
||
#[derive(Debug, Clone, Copy)]
|
||
pub(crate) struct Window {
|
||
pub(crate) x: f32,
|
||
pub(crate) y: f32,
|
||
pub(crate) w: f32,
|
||
pub(crate) h: f32,
|
||
}
|
||
|
||
/// The scale-and-pad that fits an arbitrary rectangle into the square input.
|
||
#[derive(Debug, Clone, Copy)]
|
||
pub(crate) struct Letterbox {
|
||
/// Input pixels per source pixel.
|
||
scale: f32,
|
||
pad_x: f32,
|
||
pad_y: f32,
|
||
}
|
||
|
||
impl Letterbox {
|
||
pub(crate) fn fit(w: f32, h: f32) -> Self {
|
||
let scale = (INPUT_EDGE as f32 / w).min(INPUT_EDGE as f32 / h);
|
||
Self {
|
||
scale,
|
||
pad_x: (INPUT_EDGE as f32 - w * scale) * 0.5,
|
||
pad_y: (INPUT_EDGE as f32 - h * scale) * 0.5,
|
||
}
|
||
}
|
||
|
||
/// Input pixels per source pixel.
|
||
///
|
||
/// An accessor rather than a public field so that the one place that
|
||
/// converts a model-space size into a source-space one — `Scene::cell_pixels`
|
||
/// — reads the same number this struct fitted, instead of recomputing a
|
||
/// letterbox inverse and drifting from it.
|
||
pub(crate) fn scale(&self) -> f32 {
|
||
self.scale
|
||
}
|
||
|
||
/// Resample a source window into the graph's `[1, 3, 640, 640]` input.
|
||
///
|
||
/// Bilinear, and grey (`0.5`) in the padding — the value the network sees
|
||
/// least as an edge, where black would draw a hard border across the frame
|
||
/// and invite a detection along it.
|
||
pub(crate) fn sample(
|
||
&self,
|
||
rgb: &[f32],
|
||
width: usize,
|
||
height: usize,
|
||
window: &Window,
|
||
) -> Array4<f32> {
|
||
let mut input = Array4::<f32>::from_elem((1, 3, INPUT_EDGE, INPUT_EDGE), 0.5);
|
||
|
||
for iy in 0..INPUT_EDGE {
|
||
let sy = (iy as f32 + 0.5 - self.pad_y) / self.scale + window.y;
|
||
if sy < window.y || sy >= window.y + window.h {
|
||
continue;
|
||
}
|
||
for ix in 0..INPUT_EDGE {
|
||
let sx = (ix as f32 + 0.5 - self.pad_x) / self.scale + window.x;
|
||
if sx < window.x || sx >= window.x + window.w {
|
||
continue;
|
||
}
|
||
let (x0, y0) = (sx.floor(), sy.floor());
|
||
let (fx, fy) = (sx - x0, sy - y0);
|
||
let x0 = (x0 as isize).clamp(0, width as isize - 1) as usize;
|
||
let y0 = (y0 as isize).clamp(0, height as isize - 1) as usize;
|
||
let x1 = (x0 + 1).min(width - 1);
|
||
let y1 = (y0 + 1).min(height - 1);
|
||
|
||
for c in 0..3 {
|
||
let at = |x: usize, y: usize| rgb[(y * width + x) * 3 + c];
|
||
let top = at(x0, y0) * (1.0 - fx) + at(x1, y0) * fx;
|
||
let bot = at(x0, y1) * (1.0 - fx) + at(x1, y1) * fx;
|
||
input[[0, c, iy, ix]] = top * (1.0 - fy) + bot * fy;
|
||
}
|
||
}
|
||
}
|
||
|
||
input
|
||
}
|
||
|
||
/// Input-space box back to source pixels.
|
||
fn to_source(self, x0: f32, y0: f32, x1: f32, y1: f32, w: &Window) -> (f32, f32, f32, f32) {
|
||
(
|
||
(x0 - self.pad_x) / self.scale + w.x,
|
||
(y0 - self.pad_y) / self.scale + w.y,
|
||
(x1 - self.pad_x) / self.scale + w.x,
|
||
(y1 - self.pad_y) / self.scale + w.y,
|
||
)
|
||
}
|
||
|
||
/// Source pixel to the coordinates of an output grid `stride` times
|
||
/// coarser than the graph's input.
|
||
///
|
||
/// Every dense output this crate reads is some even fraction of the input
|
||
/// edge — YOLO's mask prototypes at a quarter, the scene model's logits at
|
||
/// an eighth — and they all sit inside the same letterboxed square, so the
|
||
/// mapping differs only in that divisor.
|
||
pub(crate) fn to_grid(self, sx: f32, sy: f32, w: &Window, stride: f32) -> (f32, f32) {
|
||
(
|
||
((sx - w.x) * self.scale + self.pad_x) / stride,
|
||
((sy - w.y) * self.scale + self.pad_y) / stride,
|
||
)
|
||
}
|
||
|
||
/// Source pixel to prototype-grid coordinates.
|
||
fn to_proto(self, sx: f32, sy: f32, w: &Window) -> (f32, f32) {
|
||
self.to_grid(sx, sy, w, PROTO_STRIDE as f32)
|
||
}
|
||
}
|
||
|
||
/// Combine the prototype masks by one detection's coefficients.
|
||
///
|
||
/// The mask is `sigmoid(Σ coeff_k · proto_k)`, sampled straight into source
|
||
/// resolution and **clipped to the detection's box** — YOLO's prototypes are
|
||
/// global, so a coefficient set that describes a dog also lights up faintly on
|
||
/// a second dog elsewhere in the frame. The box is what makes an instance mask
|
||
/// an *instance* mask.
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn assemble_mask(
|
||
coeffs: &[f32],
|
||
protos: ArrayView3<f32>,
|
||
ph: usize,
|
||
pw: usize,
|
||
letterbox: &Letterbox,
|
||
window: &Window,
|
||
bbox: &(f32, f32, f32, f32),
|
||
width: usize,
|
||
height: usize,
|
||
options: &SemanticOptions,
|
||
) -> Vec<f32> {
|
||
let mut mask = vec![0.0f32; width * height];
|
||
|
||
let x0 = bbox.0.floor().max(0.0) as usize;
|
||
let y0 = bbox.1.floor().max(0.0) as usize;
|
||
let x1 = (bbox.2.ceil() as usize).min(width);
|
||
let y1 = (bbox.3.ceil() as usize).min(height);
|
||
|
||
for y in y0..y1 {
|
||
for x in x0..x1 {
|
||
let (gx, gy) = letterbox.to_proto(x as f32 + 0.5, y as f32 + 0.5, window);
|
||
if gx < 0.0 || gy < 0.0 || gx >= pw as f32 || gy >= ph as f32 {
|
||
continue;
|
||
}
|
||
|
||
// Bilinear over the prototype grid: nearest-neighbour here shows
|
||
// as visible 4-pixel stair-stepping on the mask edge.
|
||
let (fx0, fy0) = (gx.floor(), gy.floor());
|
||
let (fx, fy) = (gx - fx0, gy - fy0);
|
||
let (gx0, gy0) = (fx0 as usize, fy0 as usize);
|
||
let (gx1, gy1) = ((gx0 + 1).min(pw - 1), (gy0 + 1).min(ph - 1));
|
||
|
||
let mut acc = 0.0;
|
||
for (k, &c) in coeffs
|
||
.iter()
|
||
.enumerate()
|
||
.take(PROTOTYPES.min(protos.shape()[0]))
|
||
{
|
||
if c == 0.0 {
|
||
continue;
|
||
}
|
||
let p = protos.index_axis(ndarray::Axis(0), k);
|
||
let top = p[[gy0, gx0]] * (1.0 - fx) + p[[gy0, gx1]] * fx;
|
||
let bot = p[[gy1, gx0]] * (1.0 - fx) + p[[gy1, gx1]] * fx;
|
||
acc += c * (top * (1.0 - fy) + bot * fy);
|
||
}
|
||
|
||
let p = 1.0 / (1.0 + (-acc).exp());
|
||
if p >= options.mask_threshold * 0.5 {
|
||
mask[y * width + x] = p;
|
||
}
|
||
}
|
||
}
|
||
|
||
mask
|
||
}
|
||
|
||
/// Fold one tile's detections into the running set.
|
||
///
|
||
/// Only needed for [`Tiling::Grid`]: an object straddling a seam is seen by
|
||
/// both tiles, and without this it would appear twice in the list a person
|
||
/// chooses from.
|
||
///
|
||
/// # Why the two are unioned rather than one of them chosen
|
||
///
|
||
/// This kept the higher-scoring copy and discarded the other. That is the
|
||
/// wrong answer for the case tiling exists to serve, and it quietly threw away
|
||
/// what tiling had just paid for.
|
||
///
|
||
/// A tile sees the part of an object that falls inside it and nothing of the
|
||
/// rest, so an object on a seam produces *two partial masks*, neither of them
|
||
/// the object. Keeping the better one keeps the larger fragment — a bird with
|
||
/// its tail cut off at the tile edge, described by whichever tile held more of
|
||
/// the bird. Both halves exist; only one survived.
|
||
///
|
||
/// Unioning is sound precisely because the tiles overlap. With a 25% overlap
|
||
/// every pixel is seen by at least one tile at full resolution and pixels near
|
||
/// a seam by two, so the pointwise maximum of the two coverages is the better
|
||
/// estimate everywhere rather than a compromise: where only one tile saw a
|
||
/// pixel it is the only opinion available, and where both did, the higher
|
||
/// value comes from the tile that had more context around it. Taking a maximum
|
||
/// of soft coverage also keeps the mask soft, which is what `prior.rs` weights
|
||
/// merges by and what a mask layer's edge treatment needs.
|
||
///
|
||
/// The score is the higher of the two rather than a blend. It is shown to a
|
||
/// photographer beside the class name and means "how sure the model is this is
|
||
/// a bird"; averaging in the tile that saw only a wingtip would make a
|
||
/// confident detection look doubtful for straddling a seam.
|
||
fn merge_into(found: &mut Vec<Instance>, batch: Vec<Instance>, options: &SemanticOptions) {
|
||
for candidate in batch {
|
||
let duplicate = found.iter_mut().find(|existing| {
|
||
existing.class_id == candidate.class_id
|
||
&& existing.iou(&candidate, options.mask_threshold) >= options.merge_iou
|
||
});
|
||
|
||
match duplicate {
|
||
Some(existing) => existing.absorb(candidate),
|
||
None => found.push(candidate),
|
||
}
|
||
}
|
||
}
|
||
|
||
fn steps(extent: f32, edge: f32, stride: f32) -> usize {
|
||
if extent <= edge {
|
||
1
|
||
} else {
|
||
(((extent - edge) / stride).ceil() as usize) + 1
|
||
}
|
||
}
|
||
|
||
/// Read the class list written beside the model by `tools/export-seg-model.sh`.
|
||
///
|
||
/// A deliberately small hand-rolled reader for a flat array of strings, rather
|
||
/// than a `serde_json` dependency for one file of one shape that this
|
||
/// repository generates itself.
|
||
pub fn parse_classes(json: &str) -> Vec<Arc<str>> {
|
||
let mut out = Vec::new();
|
||
let mut chars = json.chars().peekable();
|
||
while let Some(c) = chars.next() {
|
||
if c != '"' {
|
||
continue;
|
||
}
|
||
let mut name = String::new();
|
||
while let Some(c) = chars.next() {
|
||
match c {
|
||
'"' => break,
|
||
'\\' => name.extend(chars.next()),
|
||
_ => name.push(c),
|
||
}
|
||
}
|
||
out.push(name.into());
|
||
}
|
||
out
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn classes_parse_from_the_exported_json() {
|
||
let parsed = parse_classes("[\n \"person\",\n \"bicycle\",\n \"car\"\n]");
|
||
assert_eq!(&*parsed[0], "person");
|
||
assert_eq!(&*parsed[2], "car");
|
||
assert_eq!(parsed.len(), 3);
|
||
}
|
||
|
||
#[test]
|
||
fn letterbox_round_trips_a_landscape_window() {
|
||
let window = Window {
|
||
x: 0.0,
|
||
y: 0.0,
|
||
w: 1600.0,
|
||
h: 1067.0,
|
||
};
|
||
let lb = Letterbox::fit(window.w, window.h);
|
||
|
||
// A source point maps into the input and back to where it started.
|
||
let (ix, iy) = (
|
||
(400.0 - window.x) * lb.scale + lb.pad_x,
|
||
(300.0 - window.y) * lb.scale + lb.pad_y,
|
||
);
|
||
let (sx, sy, _, _) = lb.to_source(ix, iy, 0.0, 0.0, &window);
|
||
assert!((sx - 400.0).abs() < 1e-3, "x round-trip: {sx}");
|
||
assert!((sy - 300.0).abs() < 1e-3, "y round-trip: {sy}");
|
||
}
|
||
|
||
#[test]
|
||
fn letterbox_pads_the_short_axis_only() {
|
||
let lb = Letterbox::fit(1600.0, 1067.0);
|
||
assert!(lb.pad_x.abs() < 1e-3, "wide image should not pad in x");
|
||
assert!(lb.pad_y > 100.0, "wide image should pad in y: {}", lb.pad_y);
|
||
}
|
||
|
||
/// A subject on a seam, seen in part by each of two tiles.
|
||
///
|
||
/// The failure this pins down is not a crash and not a duplicate: it is a
|
||
/// mask that looks plausible and is missing the half of the subject that
|
||
/// fell in the other tile. Keeping the higher-scoring detection produced
|
||
/// exactly that, and it is invisible unless you already know what the
|
||
/// whole subject should have been.
|
||
#[test]
|
||
fn two_tiles_seeing_one_subject_produce_the_whole_subject() {
|
||
// The left tile sees the left half strongly and nothing of the right;
|
||
// the right tile sees the right half. Together they are one bar.
|
||
let left = instance_with(
|
||
&[0.9, 0.9, 0.9, 0.8, 0.0, 0.0, 0.0, 0.0],
|
||
0.80,
|
||
(0.0, 0.0, 4.0, 1.0),
|
||
);
|
||
let right = instance_with(
|
||
&[0.0, 0.0, 0.3, 0.7, 0.9, 0.9, 0.9, 0.0],
|
||
0.60,
|
||
(2.0, 0.0, 7.0, 1.0),
|
||
);
|
||
|
||
let options = SemanticOptions {
|
||
mask_threshold: 0.5,
|
||
merge_iou: 0.1,
|
||
..SemanticOptions::default()
|
||
};
|
||
|
||
let mut found = vec![left];
|
||
merge_into(&mut found, vec![right], &options);
|
||
|
||
assert_eq!(found.len(), 1, "one subject, not two");
|
||
let m = &found[0];
|
||
|
||
// Every pixel either tile was sure about survives. Under the old
|
||
// keep-the-better-one rule, pixels 4..=6 were lost entirely.
|
||
for i in [0, 1, 2, 3, 4, 5, 6] {
|
||
assert!(
|
||
m.mask[i] >= 0.5,
|
||
"pixel {i} was seen by a tile and must survive the merge: {:?}",
|
||
m.mask
|
||
);
|
||
}
|
||
assert!(m.mask[7] < 0.5, "a pixel neither tile saw must stay out");
|
||
|
||
// Pointwise maximum, not an average: pixel 2 is 0.9 in one tile and
|
||
// 0.3 in the other, and averaging would report 0.6 — a softer edge
|
||
// than either tile actually saw.
|
||
assert!(
|
||
(m.mask[2] - 0.9).abs() < 1e-6,
|
||
"expected the max, got {}",
|
||
m.mask[2]
|
||
);
|
||
|
||
assert!(
|
||
(m.score - 0.80).abs() < 1e-6,
|
||
"the confident detection's score survives"
|
||
);
|
||
assert_eq!(m.bbox, (0.0, 0.0, 7.0, 1.0), "the box covers both halves");
|
||
}
|
||
|
||
fn instance_with(mask: &[f32], score: f32, bbox: (f32, f32, f32, f32)) -> Instance {
|
||
Instance {
|
||
class_id: 14,
|
||
class_name: "bird".into(),
|
||
score,
|
||
bbox,
|
||
mask: mask.to_vec(),
|
||
width: mask.len(),
|
||
height: 1,
|
||
}
|
||
}
|
||
|
||
/// The seam case tiling exists for, and the one it must not double-count.
|
||
#[test]
|
||
fn grid_tiling_covers_the_frame_with_overlap() {
|
||
let model_windows = |w: usize, h: usize, overlap: f32| {
|
||
// `windows` needs no session state, so exercise it through a
|
||
// stand-in rather than loading 11 MB of weights in a unit test.
|
||
let edge = (w.min(h) as f32).min(INPUT_EDGE as f32 * 1.5);
|
||
let stride = (edge * (1.0 - overlap)).max(1.0);
|
||
(steps(w as f32, edge, stride), steps(h as f32, edge, stride))
|
||
};
|
||
|
||
let (cols, rows) = model_windows(1600, 1067, 0.25);
|
||
assert!(cols >= 2, "a 1600px frame needs more than one column");
|
||
assert_eq!(rows, 2, "1067px against a 960px tile is two rows");
|
||
}
|
||
|
||
#[test]
|
||
fn a_square_frame_is_a_single_tile() {
|
||
assert_eq!(steps(640.0, 640.0, 480.0), 1);
|
||
}
|
||
|
||
#[test]
|
||
fn duplicate_detections_across_tiles_collapse_to_the_stronger() {
|
||
let opts = SemanticOptions::default();
|
||
let make = |score: f32, on: bool| Instance {
|
||
class_id: 0,
|
||
class_name: "person".into(),
|
||
score,
|
||
bbox: (0.0, 0.0, 2.0, 2.0),
|
||
mask: if on { vec![1.0; 4] } else { vec![0.0; 4] },
|
||
width: 2,
|
||
height: 2,
|
||
};
|
||
|
||
let mut found = vec![make(0.6, true)];
|
||
merge_into(&mut found, vec![make(0.9, true)], &opts);
|
||
assert_eq!(found.len(), 1, "same object seen twice is one instance");
|
||
assert_eq!(found[0].score, 0.9, "the more confident tile wins");
|
||
|
||
// A disjoint mask is a different object and must survive.
|
||
merge_into(&mut found, vec![make(0.5, false)], &opts);
|
||
assert_eq!(found.len(), 2);
|
||
}
|
||
}
|