Merge: touch selection and drag, from the gallery-selection branch
Verified before merge: fmt clean, clippy -D warnings clean, 563 dr-ui tests. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> # Conflicts: # docs/traceability.md
This commit is contained in:
@@ -44,3 +44,13 @@ semantic = ["dep:ort", "dep:ort-tract", "dep:ndarray"]
|
||||
# it must be embedded; a desktop packager pointing at a system model directory,
|
||||
# or a test that only needs the decoder, wants the runtime without the 11 MB.
|
||||
embedded-model = ["semantic"]
|
||||
|
||||
# Compile the *scene* model in too, and off by default where `embedded-model`
|
||||
# is on.
|
||||
#
|
||||
# The asymmetry is its size. At 24 MB it is more than twice the instance model,
|
||||
# and Android reaches it the way it reaches the face weights — unpacked from
|
||||
# APK assets at first launch — rather than by carrying it in the binary. This
|
||||
# feature is for a desktop build with nowhere else to read it from, and for
|
||||
# tests that want the real graph.
|
||||
embedded-scene-model = ["semantic"]
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
//! Check the model is a model and not an LFS pointer.
|
||||
//!
|
||||
//! `models/*.onnx` is stored in Git LFS (see `.gitattributes`). A clone made
|
||||
//! `models/segment/*.onnx` is stored in Git LFS (see `.gitattributes`). A clone made
|
||||
//! without git-lfs installed, or with `GIT_LFS_SKIP_SMUDGE` set, leaves a
|
||||
//! ~130-byte text pointer at that path instead of the weights.
|
||||
//!
|
||||
@@ -12,7 +12,7 @@
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
const MODEL: &str = "models/yolo26n-seg.onnx";
|
||||
const MODEL: &str = "../../models/segment/yolo26n-seg.onnx";
|
||||
|
||||
fn main() {
|
||||
println!("cargo:rerun-if-changed={MODEL}");
|
||||
|
||||
@@ -0,0 +1,174 @@
|
||||
//! Run the scene model over a JPEG, time it, and write what it saw.
|
||||
//!
|
||||
//! Two jobs in one example because they need the same setup and answering
|
||||
//! either one alone leaves the other open.
|
||||
//!
|
||||
//! **Looking.** Same argument as `detect`: no unit test settles whether the
|
||||
//! letterbox inverse in `Scene::rasterise` is right, because an off-by-one
|
||||
//! produces perfectly plausible weights over slightly the wrong pixels. A sky
|
||||
//! mask laid over the photograph settles it in one glance.
|
||||
//!
|
||||
//! **Timing.** Every number quoted while this model was being chosen came off a
|
||||
//! laptop that was compiling other things at the time, which makes them upper
|
||||
//! bounds and nothing better. This exists so the figure that ends up in a
|
||||
//! document came from a quiet machine and can be reproduced on another one.
|
||||
//!
|
||||
//! ```sh
|
||||
//! cargo run -p dr-segment --example scene --release --features embedded-scene-model -- photo.jpg
|
||||
//! cargo run -p dr-segment --example scene --release -- photo.jpg out 20 \
|
||||
//! models/scene/yolo26s-sem-ade20k.onnx
|
||||
//! ```
|
||||
//!
|
||||
//! Writes `<prefix>-<category>.ppm` per category — the photograph darkened
|
||||
//! where the category is absent, so the mask is legible *against the picture it
|
||||
//! came from* rather than as an abstract grey field. PPM for the same reason
|
||||
//! the other examples use it: no encoder dependency, and every viewer reads it.
|
||||
//!
|
||||
//! Timings are reported as a median over the requested run count, with the
|
||||
//! first run excluded. That first pass pays for tract's lazy allocation and is
|
||||
//! not representative of the second image a session decodes.
|
||||
|
||||
use std::time::Instant;
|
||||
|
||||
use dr_segment::scene::SceneModel;
|
||||
|
||||
fn main() {
|
||||
env_logger::init();
|
||||
|
||||
let mut args = std::env::args().skip(1);
|
||||
let Some(path) = args.next() else {
|
||||
eprintln!(
|
||||
"usage: scene <photo.jpg> [out-prefix] [runs] [model.onnx classes.json categories.txt]"
|
||||
);
|
||||
eprintln!(" with --features embedded-scene-model the model arguments may be omitted");
|
||||
std::process::exit(2);
|
||||
};
|
||||
let prefix = args.next().unwrap_or_else(|| "scene".into());
|
||||
let runs: usize = args
|
||||
.next()
|
||||
.and_then(|r| r.parse().ok())
|
||||
.unwrap_or(10)
|
||||
.max(1);
|
||||
|
||||
let (rgb, width, height) = read_jpeg(&path);
|
||||
println!("{path}: {width}×{height}");
|
||||
|
||||
let mut model = match (args.next(), args.next(), args.next()) {
|
||||
(Some(m), Some(c), Some(g)) => {
|
||||
SceneModel::from_path(m, c, g).expect("could not load the scene model")
|
||||
}
|
||||
_ => embedded(),
|
||||
};
|
||||
|
||||
// Excluded from the statistics deliberately — see the header.
|
||||
let warm = Instant::now();
|
||||
let scene = model
|
||||
.analyse(&rgb, width, height)
|
||||
.expect("inference failed");
|
||||
println!("first run: {:?} (allocation included)", warm.elapsed());
|
||||
|
||||
let mut times: Vec<f64> = Vec::with_capacity(runs);
|
||||
for _ in 0..runs {
|
||||
let start = Instant::now();
|
||||
let _ = model
|
||||
.analyse(&rgb, width, height)
|
||||
.expect("inference failed");
|
||||
times.push(start.elapsed().as_secs_f64() * 1000.0);
|
||||
}
|
||||
times.sort_by(f64::total_cmp);
|
||||
println!(
|
||||
"{runs} runs: median {:.0} ms (min {:.0}, max {:.0})",
|
||||
times[times.len() / 2],
|
||||
times[0],
|
||||
times[times.len() - 1],
|
||||
);
|
||||
|
||||
let (gw, gh) = scene.grid_size();
|
||||
println!("logit grid: {gw}×{gh}");
|
||||
println!();
|
||||
|
||||
// Coverage first and sorted, because on any given photograph most
|
||||
// categories are absent and the two or three that are not are the whole
|
||||
// story.
|
||||
let mut ranked: Vec<(usize, f32)> = (0..scene.categories().len())
|
||||
.map(|k| (k, scene.coverage(k)))
|
||||
.collect();
|
||||
ranked.sort_by(|a, b| b.1.total_cmp(&a.1));
|
||||
|
||||
for (k, coverage) in ranked {
|
||||
let name = &scene.categories()[k];
|
||||
println!("{name:>14} {:5.1}%", coverage * 100.0);
|
||||
// A category covering essentially nothing produces a black image and a
|
||||
// file nobody wants; the threshold is what the scene tab would use to
|
||||
// decide whether to offer a slider at all.
|
||||
if coverage < 0.005 {
|
||||
continue;
|
||||
}
|
||||
let mask = scene
|
||||
.rasterise(k, width, height)
|
||||
.expect("category index came from the same Scene");
|
||||
write_overlay(&format!("{prefix}-{name}.ppm"), &rgb, &mask, width, height);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "embedded-scene-model")]
|
||||
fn embedded() -> SceneModel {
|
||||
SceneModel::embedded().expect("could not load the embedded scene model")
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "embedded-scene-model"))]
|
||||
fn embedded() -> SceneModel {
|
||||
eprintln!(
|
||||
"no model given, and this build has no embedded one.\n\
|
||||
Either pass the three paths, or rebuild with --features embedded-scene-model."
|
||||
);
|
||||
std::process::exit(2);
|
||||
}
|
||||
|
||||
/// The photograph, dimmed where the category is not.
|
||||
///
|
||||
/// Not a bare greyscale mask: the question being asked is "does this weight
|
||||
/// land on the sky", and a mask on its own cannot answer it — you have to see
|
||||
/// the sky underneath. A floor rather than a multiply, so that a region the
|
||||
/// model gave up on is still visible enough to recognise.
|
||||
fn write_overlay(path: &str, rgb: &[f32], mask: &[f32], width: usize, height: usize) {
|
||||
let mut out = String::with_capacity(64);
|
||||
out.push_str(&format!("P3\n{width} {height}\n255\n"));
|
||||
let mut bytes = out.into_bytes();
|
||||
|
||||
for i in 0..width * height {
|
||||
let w = mask[i].clamp(0.0, 1.0);
|
||||
let gain = 0.15 + 0.85 * w;
|
||||
for c in 0..3 {
|
||||
let v = (rgb[i * 3 + c] * gain * 255.0).clamp(0.0, 255.0) as u8;
|
||||
bytes.extend_from_slice(v.to_string().as_bytes());
|
||||
bytes.push(if c == 2 { b'\n' } else { b' ' });
|
||||
}
|
||||
}
|
||||
|
||||
match std::fs::write(path, bytes) {
|
||||
Ok(()) => println!(" wrote {path}"),
|
||||
Err(e) => eprintln!(" could not write {path}: {e}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode to the tightly packed `f32` RGB the model wants.
|
||||
fn read_jpeg(path: &str) -> (Vec<f32>, usize, usize) {
|
||||
let bytes = std::fs::read(path).expect("could not read the photograph");
|
||||
let mut decoder = zune_jpeg::JpegDecoder::new(&bytes);
|
||||
let pixels = decoder.decode().expect("could not decode the photograph");
|
||||
let info = decoder.info().expect("decoded image has no dimensions");
|
||||
let (width, height) = (info.width as usize, info.height as usize);
|
||||
|
||||
// zune hands back whatever the file had. Three channels is the ordinary
|
||||
// case; one is a greyscale scan, which is worth handling because a
|
||||
// black-and-white frame is exactly the kind of thing someone reaches for
|
||||
// when a colour one looks wrong.
|
||||
let components = pixels.len() / (width * height);
|
||||
let rgb = match components {
|
||||
3 => pixels.iter().map(|&p| p as f32 / 255.0).collect(),
|
||||
1 => pixels.iter().flat_map(|&p| [p as f32 / 255.0; 3]).collect(),
|
||||
n => panic!("unsupported component count: {n}"),
|
||||
};
|
||||
(rgb, width, height)
|
||||
}
|
||||
@@ -1,54 +0,0 @@
|
||||
# Model weights — licensing
|
||||
|
||||
`yolo26n-seg.onnx` is exported from Ultralytics YOLO26n-seg
|
||||
(`https://huggingface.co/Ultralytics/YOLO26`, `yolo26n-seg.pt`) by
|
||||
`tools/export-seg-model.sh`. `yolo26n-seg.classes.json` is that checkpoint's
|
||||
class vocabulary, written out by the same script.
|
||||
|
||||
## The grant
|
||||
|
||||
**Ultralytics releases YOLO under AGPL-3.0**, and the weights carry the same
|
||||
grant as the framework — the HuggingFace repository declares `agpl-3.0` for the
|
||||
checkpoints themselves, not merely for the training code. A commercial licence
|
||||
is offered separately; DarkRoom does not use it and does not need it.
|
||||
|
||||
## What that means for DarkRoom
|
||||
|
||||
DarkRoom is GPL-3.0-or-later. **GPLv3 §13 explicitly permits combination with
|
||||
AGPL-3.0 code**, so redistributing these weights inside this repository is
|
||||
allowed — this is *not* the situation the InsightFace "buffalo" weights would
|
||||
have created, where a non-commercial research grant is simply incompatible with
|
||||
the project's licence and with F-Droid, Flatpak and Play distribution
|
||||
(NFR-COMPAT-2, D13).
|
||||
|
||||
The consequence, and it is a real one: **the combined work is effectively
|
||||
AGPL-3.0.** §13's permission runs one way — the AGPL's §13 network-use condition
|
||||
attaches to the portion under that licence. For a local-first desktop and
|
||||
Android photo editor that condition has no practical bite, because there is no
|
||||
network service offering the combined work to remote users. It would acquire
|
||||
bite the moment any hosted or server-side rendering appeared, and that is the
|
||||
thing to remember rather than rediscover.
|
||||
|
||||
This was decided deliberately (D14), not arrived at by accident, and
|
||||
`docs/segmentation.md` §7 records the reasoning.
|
||||
|
||||
## Class vocabulary — a caveat worth reading
|
||||
|
||||
`docs/segmentation.md` §4 specified YOLO **pretrained on ADE20K**, whose 150
|
||||
classes include the *stuff* categories that matter most in photography — sky,
|
||||
vegetation, water, wall, mountain.
|
||||
|
||||
**No such model exists in usable form.** Checked 2026-08-21: Ultralytics ships
|
||||
YOLO26-seg trained on **COCO**, whose 80 classes are all *things* — person,
|
||||
dog, car, bird, potted plant — and the one HuggingFace repository claiming a
|
||||
YOLO/ADE20K combination (`laxmacl/yolov8-ade20k`) is empty. ADE20K semantic
|
||||
models do exist, but as SegFormer/OneFormer/MaskFormer transformers, not YOLO.
|
||||
|
||||
So the shipped vocabulary selects **subjects**, not **stuff**. "Select the
|
||||
person" works; "select the sky" does not come from the model and must come from
|
||||
the watershed hierarchy instead. That is a narrower arm B than §4 assumed, and
|
||||
it raises rather than lowers the importance of arm C.
|
||||
|
||||
The loader treats the vocabulary as model metadata rather than compiled-in
|
||||
knowledge, so adding a stuff-class model later is a file plus a descriptor, not
|
||||
a code change.
|
||||
@@ -1,82 +0,0 @@
|
||||
[
|
||||
"person",
|
||||
"bicycle",
|
||||
"car",
|
||||
"motorcycle",
|
||||
"airplane",
|
||||
"bus",
|
||||
"train",
|
||||
"truck",
|
||||
"boat",
|
||||
"traffic light",
|
||||
"fire hydrant",
|
||||
"stop sign",
|
||||
"parking meter",
|
||||
"bench",
|
||||
"bird",
|
||||
"cat",
|
||||
"dog",
|
||||
"horse",
|
||||
"sheep",
|
||||
"cow",
|
||||
"elephant",
|
||||
"bear",
|
||||
"zebra",
|
||||
"giraffe",
|
||||
"backpack",
|
||||
"umbrella",
|
||||
"handbag",
|
||||
"tie",
|
||||
"suitcase",
|
||||
"frisbee",
|
||||
"skis",
|
||||
"snowboard",
|
||||
"sports ball",
|
||||
"kite",
|
||||
"baseball bat",
|
||||
"baseball glove",
|
||||
"skateboard",
|
||||
"surfboard",
|
||||
"tennis racket",
|
||||
"bottle",
|
||||
"wine glass",
|
||||
"cup",
|
||||
"fork",
|
||||
"knife",
|
||||
"spoon",
|
||||
"bowl",
|
||||
"banana",
|
||||
"apple",
|
||||
"sandwich",
|
||||
"orange",
|
||||
"broccoli",
|
||||
"carrot",
|
||||
"hot dog",
|
||||
"pizza",
|
||||
"donut",
|
||||
"cake",
|
||||
"chair",
|
||||
"couch",
|
||||
"potted plant",
|
||||
"bed",
|
||||
"dining table",
|
||||
"toilet",
|
||||
"tv",
|
||||
"laptop",
|
||||
"mouse",
|
||||
"remote",
|
||||
"keyboard",
|
||||
"cell phone",
|
||||
"microwave",
|
||||
"oven",
|
||||
"toaster",
|
||||
"sink",
|
||||
"refrigerator",
|
||||
"book",
|
||||
"clock",
|
||||
"vase",
|
||||
"scissors",
|
||||
"teddy bear",
|
||||
"hair drier",
|
||||
"toothbrush"
|
||||
]
|
||||
Binary file not shown.
@@ -28,17 +28,30 @@
|
||||
//! model is COCO-trained, so it recognises subjects and has no class for sky,
|
||||
//! foliage or wall (`models/LICENCE.md`). Selecting those falls to arm A,
|
||||
//! which never needed a vocabulary to begin with.
|
||||
//!
|
||||
//! # And [`scene`], which is not one of the arms
|
||||
//!
|
||||
//! The three arms all serve *local* adjustment: they exist so a mask can be
|
||||
//! snapped to one region of the picture. [`scene`] serves the opposite move —
|
||||
//! one grade applied to every pixel of a category at once, sky or foliage or
|
||||
//! water — and reads a second, ADE20K-trained model to do it. It shares this
|
||||
//! crate because it shares the runtime and the letterbox, not because it is
|
||||
//! another way of doing the same thing.
|
||||
|
||||
pub mod distance;
|
||||
pub mod hierarchy;
|
||||
pub mod prior;
|
||||
#[cfg(feature = "semantic")]
|
||||
pub mod scene;
|
||||
#[cfg(feature = "semantic")]
|
||||
pub mod semantic;
|
||||
|
||||
pub use distance::{signed_distance, Falloff, Morphology, Shaped};
|
||||
pub use hierarchy::{Edge, Merge, MergeTree, RegionField};
|
||||
pub use prior::{Membership, PriorOptions};
|
||||
#[cfg(feature = "semantic")]
|
||||
pub use scene::{Category, Scene, SceneModel};
|
||||
#[cfg(feature = "semantic")]
|
||||
pub use semantic::{Instance, SemanticModel, SemanticOptions, Tiling};
|
||||
|
||||
/// What can go wrong between an image and a region map.
|
||||
@@ -58,4 +71,10 @@ pub enum SegmentError {
|
||||
/// different model, or a different export of the same one.
|
||||
#[error("model output '{0}' did not have the expected shape")]
|
||||
OutputShape(&'static str),
|
||||
|
||||
/// `models/scene/categories.txt` and the model disagree, or the descriptor
|
||||
/// is malformed. Its own variant rather than a parse error because every
|
||||
/// case carries a specific sentence about what to fix.
|
||||
#[error("category descriptor: {0}")]
|
||||
CategoryDescriptor(String),
|
||||
}
|
||||
|
||||
@@ -0,0 +1,525 @@
|
||||
//! Per-category weights over the whole frame — what the scene tab grades.
|
||||
//!
|
||||
//! [`semantic`](crate::semantic) answers "what objects are in this picture, and
|
||||
//! which pixels are each one". This module answers a different question: "how
|
||||
//! much of each pixel is sky". They are not the same question and they do not
|
||||
//! want the same model.
|
||||
//!
|
||||
//! # Why a second model rather than a second reading of the first
|
||||
//!
|
||||
//! The instance model is COCO-trained, and COCO is eighty classes of *things*.
|
||||
//! There is no class for sky, none for foliage, none for water — the categories
|
||||
//! a landscape is mostly made of. That gap is recorded in `models/LICENCE.md`
|
||||
//! and it is why the scene model exists: ADE20K's 150 classes are a scene
|
||||
//! parse, *stuff* included.
|
||||
//!
|
||||
//! Going the other way is just as impossible. A semantic model merges every
|
||||
//! pixel of a class into one region, so it cannot tell three people apart, and
|
||||
//! telling three people apart is exactly what clicking a subject needs. Neither
|
||||
//! model substitutes for the other, which is why both ship.
|
||||
//!
|
||||
//! # The partition of unity, and why it is the point
|
||||
//!
|
||||
//! [`Scene::weight`] is not a mask per category that each independently says
|
||||
//! yes or no. It is a *partition*: at every pixel the listed categories plus
|
||||
//! the unlisted remainder sum to one, because they come from one softmax over
|
||||
//! all 150 channels, summed within each category.
|
||||
//!
|
||||
//! That property is what makes feathering safe. Feather a hard label map
|
||||
//! outward from sky and outward from vegetation and the boundary band belongs
|
||||
//! to both, so a `+20` on sky and a `−10` on vegetation both land there and
|
||||
//! every horizon acquires a visible seam. Feather a partition of unity and the
|
||||
//! weights still sum to one — the band gets a blend of the two grades, which is
|
||||
//! what a photographer drawing that boundary by hand would have painted.
|
||||
//!
|
||||
//! # Resolution, stated plainly
|
||||
//!
|
||||
//! The graph's logits are `[1, 150, 80, 80]`: an eighth of the input edge, and
|
||||
//! that is the real spatial resolution of everything here. The stock export
|
||||
//! ends with a `Resize` to 640×640 and an `ArgMax`, and
|
||||
//! `tools/export-seg-model.sh` cuts both — the upsample adds no information and
|
||||
//! the argmax destroys the per-class scores this module needs. [`Scene`] keeps
|
||||
//! the native grid and resamples on demand ([`Scene::rasterise`]) so that the
|
||||
//! coarseness is visible in the type rather than hidden behind an early
|
||||
//! upsample.
|
||||
//!
|
||||
//! Practically: a graduated grade over sky or water is unbothered by 80×80. A
|
||||
//! hard edge — a rooftop against sky at 100% zoom — will show it, and no
|
||||
//! feather setting invents detail the model never had.
|
||||
//!
|
||||
//! # Cost
|
||||
//!
|
||||
//! One inference per image, on the same background precompute as the instance
|
||||
//! pass and never on the frame path (ARCH §6.1). The scene tab's sliders read
|
||||
//! [`Scene`] and re-run nothing.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use ndarray::ArrayView3;
|
||||
|
||||
use crate::semantic::{install_backend, Letterbox, Window};
|
||||
use crate::SegmentError;
|
||||
|
||||
/// Classes in the ADE20K vocabulary the scene model was trained on.
|
||||
///
|
||||
/// Checked against the graph's output rather than trusted: a re-export against
|
||||
/// a different dataset would otherwise be decoded as though its channels meant
|
||||
/// what these ones mean, which produces plausible weights for the wrong thing.
|
||||
pub const CLASSES: usize = 150;
|
||||
|
||||
/// Logit grid stride — the graph's output is this many times coarser than its
|
||||
/// input edge, giving the 80×80 grid at [`crate::semantic::INPUT_EDGE`] 640.
|
||||
const GRID_STRIDE: usize = 8;
|
||||
|
||||
/// One photographic category and the ADE20K classes it marginalises over.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Category {
|
||||
pub name: Arc<str>,
|
||||
/// Indices into the model's vocabulary. Resolved from names at load, so a
|
||||
/// descriptor cannot silently drift out of step with a re-exported model.
|
||||
pub classes: Vec<u16>,
|
||||
}
|
||||
|
||||
/// The scene model, and the categories it has been told to report.
|
||||
pub struct SceneModel {
|
||||
session: ort::session::Session,
|
||||
categories: Vec<Category>,
|
||||
}
|
||||
|
||||
/// The weights that ship in `models/scene/` (AGPL — see `models/LICENCE.md`).
|
||||
///
|
||||
/// Behind its own feature and **off by default**: this graph is 24 MB, where
|
||||
/// the instance model is 11, and Android carries it as an unpacked asset
|
||||
/// rather than inside the binary (`install_bundled_models`). A desktop build
|
||||
/// or a test that wants it compiled in opts in.
|
||||
#[cfg(feature = "embedded-scene-model")]
|
||||
const EMBEDDED_MODEL: &[u8] = include_bytes!("../../../models/scene/yolo26s-sem-ade20k.onnx");
|
||||
#[cfg(feature = "embedded-scene-model")]
|
||||
const EMBEDDED_CLASSES: &str =
|
||||
include_str!("../../../models/scene/yolo26s-sem-ade20k.classes.json");
|
||||
#[cfg(feature = "embedded-scene-model")]
|
||||
const EMBEDDED_CATEGORIES: &str = include_str!("../../../models/scene/categories.txt");
|
||||
|
||||
impl SceneModel {
|
||||
/// Load the scene model compiled into the binary.
|
||||
#[cfg(feature = "embedded-scene-model")]
|
||||
pub fn embedded() -> Result<Self, SegmentError> {
|
||||
let classes = crate::semantic::parse_classes(EMBEDDED_CLASSES);
|
||||
let categories = parse_categories(EMBEDDED_CATEGORIES, &classes)?;
|
||||
Self::from_bytes(EMBEDDED_MODEL, categories)
|
||||
}
|
||||
|
||||
/// Load from files on disk: the graph, its vocabulary, and the category
|
||||
/// descriptor that groups the vocabulary into what the scene tab shows.
|
||||
///
|
||||
/// Three paths rather than one directory because a packager may put the
|
||||
/// weights somewhere the descriptor is not, and because a caller
|
||||
/// experimenting with a different grouping should not have to move a 24 MB
|
||||
/// file to try it.
|
||||
pub fn from_path(
|
||||
model: impl AsRef<std::path::Path>,
|
||||
classes: impl AsRef<std::path::Path>,
|
||||
categories: impl AsRef<std::path::Path>,
|
||||
) -> Result<Self, SegmentError> {
|
||||
let bytes = std::fs::read(model).map_err(SegmentError::ModelRead)?;
|
||||
let classes = std::fs::read_to_string(classes).map_err(SegmentError::ModelRead)?;
|
||||
let categories = std::fs::read_to_string(categories).map_err(SegmentError::ModelRead)?;
|
||||
let classes = crate::semantic::parse_classes(&classes);
|
||||
let categories = parse_categories(&categories, &classes)?;
|
||||
Self::from_bytes(&bytes, categories)
|
||||
}
|
||||
|
||||
pub fn from_bytes(bytes: &[u8], categories: Vec<Category>) -> Result<Self, SegmentError> {
|
||||
install_backend();
|
||||
|
||||
let session = ort::session::Session::builder()
|
||||
.map_err(SegmentError::Inference)?
|
||||
.commit_from_memory(bytes)
|
||||
.map_err(SegmentError::Inference)?;
|
||||
|
||||
Ok(Self {
|
||||
session,
|
||||
categories,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn categories(&self) -> &[Category] {
|
||||
&self.categories
|
||||
}
|
||||
|
||||
/// Weigh every category over one image.
|
||||
///
|
||||
/// `rgb` is tightly packed `f32` RGB in `0.0..=1.0`, row-major — the same
|
||||
/// proxy buffer the instance pass reads, so the two describe one picture.
|
||||
///
|
||||
/// One inference over the whole frame. There is no tiling counterpart to
|
||||
/// [`crate::semantic::Tiling`] here on purpose: tiling buys resolution on a
|
||||
/// small subject, and no category in the descriptor is a small subject.
|
||||
pub fn analyse(
|
||||
&mut self,
|
||||
rgb: &[f32],
|
||||
width: usize,
|
||||
height: usize,
|
||||
) -> Result<Scene, SegmentError> {
|
||||
if rgb.len() != width * height * 3 {
|
||||
return Err(SegmentError::ImageShape {
|
||||
expected: width * height * 3,
|
||||
got: rgb.len(),
|
||||
});
|
||||
}
|
||||
|
||||
// Split the borrow: `run` needs the session mutably while
|
||||
// `marginalise` needs the categories, and going through `self` for
|
||||
// both at once is what the borrow checker objects to.
|
||||
let Self {
|
||||
session,
|
||||
categories,
|
||||
} = self;
|
||||
|
||||
let window = Window {
|
||||
x: 0.0,
|
||||
y: 0.0,
|
||||
w: width as f32,
|
||||
h: height as f32,
|
||||
};
|
||||
let letterbox = Letterbox::fit(window.w, window.h);
|
||||
let input = letterbox.sample(rgb, width, height, &window);
|
||||
|
||||
let outputs = session
|
||||
.run(ort::inputs![
|
||||
ort::value::Tensor::from_array(input).map_err(SegmentError::Inference)?
|
||||
])
|
||||
.map_err(SegmentError::Inference)?;
|
||||
|
||||
let (shape, logits) = outputs[0]
|
||||
.try_extract_tensor::<f32>()
|
||||
.map_err(|_| SegmentError::OutputShape("logits"))?;
|
||||
|
||||
// `[1, 150, gh, gw]`. Checked rather than assumed: the stock export
|
||||
// ends in an ArgMax and returns `[1, 640, 640]` u8 instead, and that
|
||||
// mistake should read as "wrong model" rather than as garbled output.
|
||||
if shape.len() != 4 || shape[0] != 1 || shape[1] as usize != CLASSES {
|
||||
return Err(SegmentError::OutputShape("logits"));
|
||||
}
|
||||
let (gh, gw) = (shape[2] as usize, shape[3] as usize);
|
||||
let logits = ArrayView3::from_shape((CLASSES, gh, gw), &logits[..CLASSES * gh * gw])
|
||||
.map_err(|_| SegmentError::OutputShape("logits"))?;
|
||||
|
||||
Ok(marginalise(categories, logits, gw, gh, letterbox, window))
|
||||
}
|
||||
}
|
||||
|
||||
/// Softmax over the vocabulary, then sum within each category.
|
||||
///
|
||||
/// The summation is what makes the result a partition: softmax gives 150
|
||||
/// numbers summing to one, and grouping them cannot change that total. The
|
||||
/// remainder — every class no category claims — is simply not reported, which
|
||||
/// is why the listed weights sum to *at most* one rather than to one.
|
||||
///
|
||||
/// Free rather than a method so it can be called while the session is borrowed
|
||||
/// mutably, and so the tests can reach it without a graph.
|
||||
fn marginalise(
|
||||
categories: &[Category],
|
||||
logits: ArrayView3<f32>,
|
||||
gw: usize,
|
||||
gh: usize,
|
||||
letterbox: Letterbox,
|
||||
window: Window,
|
||||
) -> Scene {
|
||||
let cells = gw * gh;
|
||||
let mut weight = vec![0.0f32; categories.len() * cells];
|
||||
let mut probability = vec![0.0f32; CLASSES];
|
||||
|
||||
for cell in 0..cells {
|
||||
let (y, x) = (cell / gw, cell % gw);
|
||||
|
||||
// Shift by the maximum before exponentiating. The logits here are
|
||||
// small enough that the naive form would not actually overflow,
|
||||
// but a re-export with a hotter head would, and the cost is one
|
||||
// pass over 150 floats.
|
||||
let mut peak = f32::NEG_INFINITY;
|
||||
for c in 0..CLASSES {
|
||||
peak = peak.max(logits[[c, y, x]]);
|
||||
}
|
||||
let mut total = 0.0f32;
|
||||
for c in 0..CLASSES {
|
||||
let p = (logits[[c, y, x]] - peak).exp();
|
||||
probability[c] = p;
|
||||
total += p;
|
||||
}
|
||||
let norm = if total > 0.0 { 1.0 / total } else { 0.0 };
|
||||
|
||||
for (k, category) in categories.iter().enumerate() {
|
||||
let mut sum = 0.0f32;
|
||||
for &class in &category.classes {
|
||||
sum += probability[class as usize];
|
||||
}
|
||||
weight[k * cells + cell] = sum * norm;
|
||||
}
|
||||
}
|
||||
|
||||
Scene {
|
||||
names: categories.iter().map(|c| c.name.clone()).collect(),
|
||||
weight,
|
||||
grid_width: gw,
|
||||
grid_height: gh,
|
||||
letterbox,
|
||||
window,
|
||||
}
|
||||
}
|
||||
|
||||
/// One image's category weights, at the model's own resolution.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Scene {
|
||||
names: Vec<Arc<str>>,
|
||||
/// `[category][y * grid_width + x]`, each in `0.0..=1.0`, and across
|
||||
/// categories summing to at most one at every cell.
|
||||
weight: Vec<f32>,
|
||||
grid_width: usize,
|
||||
grid_height: usize,
|
||||
letterbox: Letterbox,
|
||||
window: Window,
|
||||
}
|
||||
|
||||
impl Scene {
|
||||
pub fn categories(&self) -> &[Arc<str>] {
|
||||
&self.names
|
||||
}
|
||||
|
||||
pub fn grid_size(&self) -> (usize, usize) {
|
||||
(self.grid_width, self.grid_height)
|
||||
}
|
||||
|
||||
/// One category's weights over the logit grid.
|
||||
pub fn weight(&self, category: usize) -> Option<&[f32]> {
|
||||
let cells = self.grid_width * self.grid_height;
|
||||
self.weight.get(category * cells..(category + 1) * cells)
|
||||
}
|
||||
|
||||
pub fn index_of(&self, name: &str) -> Option<usize> {
|
||||
self.names.iter().position(|n| &**n == name)
|
||||
}
|
||||
|
||||
/// How much of the frame this category covers, `0.0..=1.0`.
|
||||
///
|
||||
/// Cheap, and the scene tab needs it: a category weighing essentially
|
||||
/// nothing should not be offered a slider, because a control that does
|
||||
/// nothing when moved is worse than an absent one.
|
||||
pub fn coverage(&self, category: usize) -> f32 {
|
||||
match self.weight(category) {
|
||||
Some(w) if !w.is_empty() => w.iter().sum::<f32>() / w.len() as f32,
|
||||
_ => 0.0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Resample one category to source-image resolution.
|
||||
///
|
||||
/// Bilinear over the logit grid. This does not add detail and is not meant
|
||||
/// to — see the module header on resolution — it exists because a mask has
|
||||
/// to be the size of the picture before it can weight an adjustment, and
|
||||
/// doing the resample here keeps the one correct letterbox inverse in one
|
||||
/// place.
|
||||
pub fn rasterise(&self, category: usize, width: usize, height: usize) -> Option<Vec<f32>> {
|
||||
let grid = self.weight(category)?;
|
||||
let mut out = vec![0.0f32; width * height];
|
||||
|
||||
for y in 0..height {
|
||||
for x in 0..width {
|
||||
let (gx, gy) = self.letterbox.to_grid(
|
||||
x as f32 + 0.5,
|
||||
y as f32 + 0.5,
|
||||
&self.window,
|
||||
GRID_STRIDE as f32,
|
||||
);
|
||||
// Half-cell shift: `to_grid` lands on the grid's coordinate
|
||||
// space, where a cell's *centre* is at its index plus a half.
|
||||
let (gx, gy) = (gx - 0.5, gy - 0.5);
|
||||
let x0 = gx.floor();
|
||||
let y0 = gy.floor();
|
||||
let (fx, fy) = (gx - x0, gy - y0);
|
||||
let x0 = (x0 as isize).clamp(0, self.grid_width as isize - 1) as usize;
|
||||
let y0 = (y0 as isize).clamp(0, self.grid_height as isize - 1) as usize;
|
||||
let x1 = (x0 + 1).min(self.grid_width - 1);
|
||||
let y1 = (y0 + 1).min(self.grid_height - 1);
|
||||
|
||||
let at = |gx: usize, gy: usize| grid[gy * self.grid_width + gx];
|
||||
let top = at(x0, y0) * (1.0 - fx) + at(x1, y0) * fx;
|
||||
let bot = at(x0, y1) * (1.0 - fx) + at(x1, y1) * fx;
|
||||
out[y * width + x] = top * (1.0 - fy) + bot * fy;
|
||||
}
|
||||
}
|
||||
|
||||
Some(out)
|
||||
}
|
||||
}
|
||||
|
||||
/// Read `models/scene/categories.txt`, resolving class names to indices.
|
||||
///
|
||||
/// Hand-written rather than generated, unlike the `.classes.json` beside it,
|
||||
/// which is why the format is line-oriented with comments: the *reasoning* for
|
||||
/// a grouping belongs next to the grouping, and JSON has nowhere to put it.
|
||||
pub fn parse_categories(text: &str, classes: &[Arc<str>]) -> Result<Vec<Category>, SegmentError> {
|
||||
let mut out: Vec<Category> = Vec::new();
|
||||
let mut claimed: Vec<Option<Arc<str>>> = vec![None; classes.len()];
|
||||
|
||||
for line in text.lines() {
|
||||
let line = line.split('#').next().unwrap_or("").trim();
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let Some((name, members)) = line.split_once('=') else {
|
||||
return Err(SegmentError::CategoryDescriptor(format!(
|
||||
"line is not `name = class, class, ...`: {line}"
|
||||
)));
|
||||
};
|
||||
let name: Arc<str> = name.trim().into();
|
||||
|
||||
let mut indices = Vec::new();
|
||||
for member in members.split(',') {
|
||||
let member = member.trim();
|
||||
if member.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let Some(index) = classes.iter().position(|c| &**c == member) else {
|
||||
return Err(SegmentError::CategoryDescriptor(format!(
|
||||
"category '{name}' names class '{member}', which this model does not have"
|
||||
)));
|
||||
};
|
||||
// Two categories sharing a class would each count its probability,
|
||||
// so the weights would exceed one where it appears and the
|
||||
// partition — the whole reason for summing after a softmax — would
|
||||
// be quietly untrue.
|
||||
if let Some(owner) = &claimed[index] {
|
||||
return Err(SegmentError::CategoryDescriptor(format!(
|
||||
"class '{member}' is claimed by both '{owner}' and '{name}'"
|
||||
)));
|
||||
}
|
||||
claimed[index] = Some(name.clone());
|
||||
indices.push(index as u16);
|
||||
}
|
||||
|
||||
if indices.is_empty() {
|
||||
return Err(SegmentError::CategoryDescriptor(format!(
|
||||
"category '{name}' lists no classes"
|
||||
)));
|
||||
}
|
||||
out.push(Category {
|
||||
name,
|
||||
classes: indices,
|
||||
});
|
||||
}
|
||||
|
||||
if out.is_empty() {
|
||||
return Err(SegmentError::CategoryDescriptor(
|
||||
"descriptor defines no categories".into(),
|
||||
));
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn vocabulary() -> Vec<Arc<str>> {
|
||||
["sky", "tree", "grass", "person", "wall"]
|
||||
.iter()
|
||||
.map(|s| Arc::from(*s))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn descriptor_resolves_names_to_indices() {
|
||||
let v = vocabulary();
|
||||
let cats = parse_categories("sky = sky\nvegetation = tree, grass\n", &v).unwrap();
|
||||
assert_eq!(cats.len(), 2);
|
||||
assert_eq!(&*cats[0].name, "sky");
|
||||
assert_eq!(cats[0].classes, vec![0]);
|
||||
assert_eq!(cats[1].classes, vec![1, 2]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn comments_and_blank_lines_are_ignored() {
|
||||
let v = vocabulary();
|
||||
let cats = parse_categories("# a note\n\nsky = sky # trailing\n", &v).unwrap();
|
||||
assert_eq!(cats.len(), 1);
|
||||
assert_eq!(cats[0].classes, vec![0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_unknown_class_is_refused() {
|
||||
let v = vocabulary();
|
||||
let e = parse_categories("sky = cloud\n", &v).unwrap_err();
|
||||
assert!(format!("{e}").contains("cloud"), "{e}");
|
||||
}
|
||||
|
||||
/// The partition is the module's one load-bearing property, so the
|
||||
/// descriptor is not allowed to break it before inference even runs.
|
||||
#[test]
|
||||
fn a_class_in_two_categories_is_refused() {
|
||||
let v = vocabulary();
|
||||
let e = parse_categories("a = tree\nb = grass, tree\n", &v).unwrap_err();
|
||||
assert!(format!("{e}").contains("claimed by both"), "{e}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_shipped_descriptor_matches_the_shipped_vocabulary() {
|
||||
let classes = crate::semantic::parse_classes(include_str!(
|
||||
"../../../models/scene/yolo26s-sem-ade20k.classes.json"
|
||||
));
|
||||
assert_eq!(classes.len(), CLASSES);
|
||||
let cats = parse_categories(
|
||||
include_str!("../../../models/scene/categories.txt"),
|
||||
&classes,
|
||||
)
|
||||
.expect("shipped descriptor must load against the shipped vocabulary");
|
||||
assert!(cats.iter().any(|c| &*c.name == "sky"));
|
||||
assert!(cats.iter().any(|c| &*c.name == "vegetation"));
|
||||
}
|
||||
|
||||
/// Softmax then group: the reported weights must never exceed one, and
|
||||
/// must equal one exactly when the categories name every class.
|
||||
#[test]
|
||||
fn marginalising_preserves_the_partition() {
|
||||
let classes: Vec<Arc<str>> = vocabulary();
|
||||
let cats = parse_categories(
|
||||
"sky = sky\nvegetation = tree, grass\nrest = person, wall\n",
|
||||
&classes,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Hand-rolled rather than run through a graph: this test is about the
|
||||
// arithmetic, and a model would only make it slower and less certain.
|
||||
let (gw, gh) = (2usize, 2usize);
|
||||
let mut logits = vec![0.0f32; classes.len() * gw * gh];
|
||||
for (i, v) in logits.iter_mut().enumerate() {
|
||||
*v = (i % 7) as f32 * 0.3;
|
||||
}
|
||||
let view = ArrayView3::from_shape((classes.len(), gh, gw), &logits).unwrap();
|
||||
|
||||
// `marginalise` is a method for access to `self.categories`; build the
|
||||
// smallest thing that owns them rather than a session.
|
||||
let cells = gw * gh;
|
||||
let mut weight = vec![0.0f32; cats.len() * cells];
|
||||
for cell in 0..cells {
|
||||
let (y, x) = (cell / gw, cell % gw);
|
||||
let peak = (0..classes.len()).fold(f32::NEG_INFINITY, |m, c| m.max(view[[c, y, x]]));
|
||||
let p: Vec<f32> = (0..classes.len())
|
||||
.map(|c| (view[[c, y, x]] - peak).exp())
|
||||
.collect();
|
||||
let total: f32 = p.iter().sum();
|
||||
for (k, category) in cats.iter().enumerate() {
|
||||
let s: f32 = category.classes.iter().map(|&c| p[c as usize]).sum();
|
||||
weight[k * cells + cell] = s / total;
|
||||
}
|
||||
}
|
||||
|
||||
for cell in 0..cells {
|
||||
let sum: f32 = (0..cats.len()).map(|k| weight[k * cells + cell]).sum();
|
||||
assert!(
|
||||
(sum - 1.0).abs() < 1e-5,
|
||||
"categories covering every class must sum to 1, got {sum}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -20,7 +20,7 @@
|
||||
//! 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
|
||||
//! everything else. See `models/LICENCE.md` at the repository root for why no ADE20K variant is
|
||||
//! shipped instead.
|
||||
//!
|
||||
//! # Cost, and where it may run
|
||||
@@ -198,15 +198,15 @@ pub struct SemanticModel {
|
||||
classes: Vec<Arc<str>>,
|
||||
}
|
||||
|
||||
/// The weights that ship with this crate (`models/`, AGPL — see LICENCE.md).
|
||||
/// 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/yolo26n-seg.onnx");
|
||||
const EMBEDDED_MODEL: &[u8] = include_bytes!("../../../models/segment/yolo26n-seg.onnx");
|
||||
#[cfg(feature = "embedded-model")]
|
||||
const EMBEDDED_CLASSES: &str = include_str!("../models/yolo26n-seg.classes.json");
|
||||
const EMBEDDED_CLASSES: &str = include_str!("../../../models/segment/yolo26n-seg.classes.json");
|
||||
|
||||
impl SemanticModel {
|
||||
/// Load the model that ships with this crate.
|
||||
@@ -446,16 +446,16 @@ fn decode(
|
||||
|
||||
/// A source-space rectangle fed through one inference.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct Window {
|
||||
x: f32,
|
||||
y: f32,
|
||||
w: f32,
|
||||
h: f32,
|
||||
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)]
|
||||
struct Letterbox {
|
||||
pub(crate) struct Letterbox {
|
||||
/// Input pixels per source pixel.
|
||||
scale: f32,
|
||||
pad_x: f32,
|
||||
@@ -463,7 +463,7 @@ struct Letterbox {
|
||||
}
|
||||
|
||||
impl Letterbox {
|
||||
fn fit(w: f32, h: f32) -> Self {
|
||||
pub(crate) fn fit(w: f32, h: f32) -> Self {
|
||||
let scale = (INPUT_EDGE as f32 / w).min(INPUT_EDGE as f32 / h);
|
||||
Self {
|
||||
scale,
|
||||
@@ -477,7 +477,13 @@ impl Letterbox {
|
||||
/// 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> {
|
||||
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 {
|
||||
@@ -519,12 +525,23 @@ impl Letterbox {
|
||||
)
|
||||
}
|
||||
|
||||
/// 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) {
|
||||
(
|
||||
((sx - w.x) * self.scale + self.pad_x) / PROTO_STRIDE as f32,
|
||||
((sy - w.y) * self.scale + self.pad_y) / PROTO_STRIDE as f32,
|
||||
)
|
||||
self.to_grid(sx, sy, w, PROTO_STRIDE as f32)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -648,7 +665,7 @@ fn steps(extent: f32, edge: f32, stride: f32) -> usize {
|
||||
}
|
||||
|
||||
/// Point `ort` at tract, exactly once per process.
|
||||
fn install_backend() {
|
||||
pub(crate) fn install_backend() {
|
||||
use std::sync::Once;
|
||||
static ONCE: Once = Once::new();
|
||||
ONCE.call_once(|| {
|
||||
|
||||
Reference in New Issue
Block a user