dtourolleandClaude Opus 5
56978fdf35
🐳 Android image / Build and push (push) Successful in 1s
Build and test / android-image (push) Successful in 1s
Build and test / Desktop (Linux) (push) Failing after 9m6s
Build and test / Layer separation (push) Successful in 26s
Traceability / Requirement traces (push) Failing after 23s
Build and test / Android (aarch64) (push) Failing after 22m38s
Nothing here is film simulation. These are lints that fail master today,
under the -D warnings CI runs with, mostly from a toolchain that learned
new ones rather than from anybody's code -- `is_multiple_of` and the
derivable `Default` did not exist as lints when this was written.
They are fixed rather than allowed, and by hand rather than by trusting
`cargo clippy --fix` wholesale: its automatic pass split a derive in two
and left a stray blank line, which is the sort of thing that is correct
and still wrong to commit.
The four that needed a decision rather than a rewrite:
- The distance transform's inner loop writes through its iterator now.
`q` stays, because it is the position the parabola is evaluated at as
well as the index it is written to -- the lint is about the write.
- `to_source` and `to_proto` take `self` by value. Their receiver is
`Copy`, so this is the same machine code and the honest signature.
- The export path's return type is five levels deep and now has a name,
plus a line saying why the `Option` wraps the `Result`: `None` is
cancellation, which is not a failure and has no error to report.
- A test fills a range instead of looping over one.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
843 lines
31 KiB
Rust
843 lines
31 KiB
Rust
//! Semantic segmentation — arm B (S15, docs/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` 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: ort::session::Session,
|
||
classes: Vec<Arc<str>>,
|
||
}
|
||
|
||
/// The weights that ship with this crate (`models/`, AGPL — see 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/yolo26n-seg.onnx");
|
||
#[cfg(feature = "embedded-model")]
|
||
const EMBEDDED_CLASSES: &str = include_str!("../models/yolo26n-seg.classes.json");
|
||
|
||
impl SemanticModel {
|
||
/// Load the model that ships with this crate.
|
||
#[cfg(feature = "embedded-model")]
|
||
pub fn embedded() -> Result<Self, SegmentError> {
|
||
Self::from_bytes(EMBEDDED_MODEL, 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> {
|
||
// Idempotent, and it must happen before any other `ort` call: with
|
||
// `alternative-backend` there is no linked runtime to fall back on, so
|
||
// an un-set API is a panic rather than a slow path.
|
||
install_backend();
|
||
|
||
let session = ort::session::Session::builder()
|
||
.map_err(SegmentError::Inference)?
|
||
.commit_from_memory(bytes)
|
||
.map_err(SegmentError::Inference)?;
|
||
|
||
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 outputs = self
|
||
.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)]
|
||
struct Window {
|
||
x: f32,
|
||
y: f32,
|
||
w: f32,
|
||
h: f32,
|
||
}
|
||
|
||
/// The scale-and-pad that fits an arbitrary rectangle into the square input.
|
||
#[derive(Debug, Clone, Copy)]
|
||
struct Letterbox {
|
||
/// Input pixels per source pixel.
|
||
scale: f32,
|
||
pad_x: f32,
|
||
pad_y: f32,
|
||
}
|
||
|
||
impl Letterbox {
|
||
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,
|
||
}
|
||
}
|
||
|
||
/// 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.
|
||
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 prototype-grid coordinates.
|
||
fn to_proto(self, sx: f32, sy: f32, w: &Window) -> (f32, f32) {
|
||
(
|
||
((sx - w.x) * self.scale + self.pad_x) / PROTO_STRIDE as f32,
|
||
((sy - w.y) * self.scale + self.pad_y) / 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
|
||
}
|
||
}
|
||
|
||
/// Point `ort` at tract, exactly once per process.
|
||
fn install_backend() {
|
||
use std::sync::Once;
|
||
static ONCE: Once = Once::new();
|
||
ONCE.call_once(|| {
|
||
// Returns false if an API was already installed, which is not an error
|
||
// — it means something else got here first, and there is only one
|
||
// backend compiled in for it to have chosen.
|
||
let _ = ort::set_api(ort_tract::api());
|
||
});
|
||
}
|
||
|
||
/// 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);
|
||
}
|
||
}
|