Let the model say what a thing is and the watershed say where it ends
Local masking needs to know where an image's regions are. The watershed spike (S15 arm A) found the boundaries but had no idea what any of them enclosed; its coarse levels were geometric accidents. This adds the other half and the thing that joins them. `core/dr-segment` is where region reasoning now lives — the hierarchy moves out of `dr-gpu`, which keeps only the pixel passes that are genuinely shaders. The new crate is device-free and, without its default features, model-free too: 20 of its tests need neither an adapter nor 11 MB of weights. Arm B runs YOLO26n-seg through `ort`. D13 framed inference as a choice between `ort`'s C++ runtime and the pure-Rust dependency policy; that was a false choice. `ort`'s `alternative-backend` feature unlinks the C entirely and `ort-tract` supplies the API from tract, which is pure Rust. Measured before committing to it: zero unsupported operators, 420 ms for 640x640, and correct masks on bus.jpg. No NDK problem to solve, so D13's largest tolerated exception is not needed. Arm C is `prior.rs`, and it ships because the two arms fail in opposite directions. Instance membership re-weights the merge saddles, so region pairs the model believes share an object merge early and pairs straddling its edge merge late. No boundary moves — only the order in which they dissolve — which is how the result stays pixel-accurate at every level while its coarse levels become named things. Two things the spec assumed that turned out to be false, both recorded in models/LICENCE.md: there is no usable ADE20K-trained YOLO, so the shipped vocabulary is COCO's 80 subjects and *stuff* like sky and foliage must come from arm A; and tract cannot parse a dynamic-shape export, so the graph's input is fixed and tiling is the only route to more semantic resolution. Weights are AGPL-3.0, which GPLv3 §13 permits and which makes the combined work effectively AGPL. Deliberate, not accidental. They live in Git LFS, and a build script fails with an instruction rather than embedding a pointer file when the clone lacks them.
This commit is contained in:
@@ -0,0 +1,709 @@
|
||||
//! 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)]
|
||||
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.
|
||||
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 },
|
||||
}
|
||||
|
||||
impl Default for Tiling {
|
||||
fn default() -> Self {
|
||||
Self::Whole
|
||||
}
|
||||
}
|
||||
|
||||
/// 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
|
||||
}
|
||||
}
|
||||
|
||||
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. Keeps the higher-scoring copy, which is generally the tile
|
||||
/// that saw more of the object.
|
||||
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) if existing.score < candidate.score => *existing = candidate,
|
||||
Some(_) => {}
|
||||
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);
|
||||
}
|
||||
|
||||
/// 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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user