From 90a248edad7c61b272014c86be8788d0b8e8c722 Mon Sep 17 00:00:00 2001 From: helgehelge123 <117832631+helgehelge123@users.noreply.github.com> Date: Fri, 16 Oct 2026 01:35:59 +0200 Subject: [PATCH] M12: Describe selections from SAM 3's semantic mask: sure inside, no lines between instances Describe merged the detector's per-instance masks. Those stay near 0 even well inside what they found: a sky came out about 87 % selected, a volleyball net about 70 %, a person about 94 %. Two instances that touch both fade out where they meet, so a line showed between people standing side by side. - The detector's semantic head is now ported: semantic_projection over the pixel decoder, the transformers implementation's semantic_seg. It matches the reference to about 1e-3, like the instance masks. tools/sam3_reference.py writes it and the reference test checks it; older reference files still pass. - Each instance found (score above the threshold, as before) takes the semantic logits where it is the strongest instance and near (above -4). The semantic mask is sure inside and runs on across the boundary between two instances. Far from every instance nothing changes, so Describe selects nothing it didn't find before. On 16 prompts over three photos (a mountain landscape, two gym photos), the median coverage inside each selection rises, e.g.: - sky 0.87 -> 1.00 - mountains 0.81 -> 0.99 - net 0.71 -> 0.99 - floor 0.92 -> 1.00 - person 0.94 -> 0.99 Each selection gains 0 to 1.9 % of the photo at its soft rim and loses none. Co-Authored-By: Claude Opus 5.5 --- crates/segment/src/detector.rs | 98 ++++++++++++++++++++++++++++--- crates/segment/tests/reference.rs | 7 +++ docs/ai-masks.md | 7 ++- tools/sam3_reference.py | 1 + 4 files changed, 102 insertions(+), 11 deletions(-) diff --git a/crates/segment/src/detector.rs b/crates/segment/src/detector.rs index 7c4eccbd3..c302aafe7 100644 --- a/crates/segment/src/detector.rs +++ b/crates/segment/src/detector.rs @@ -41,6 +41,9 @@ pub struct Output { pub boxes: Vec<[f32; 4]>, /// The text features `[1, 32, 256]` (tests). pub text: Tensor, + /// Semantic logits `[288, 288]` flattened: everything the prompt names, one mask without + /// boundaries between instances. + pub semantic: Vec, } struct EncoderLayer { @@ -91,6 +94,7 @@ pub struct Detector { pixel_convs: Vec, pixel_norms: Vec, instance_proj: Conv2d, + semantic_proj: Conv2d, mask_embedder: Vec, } @@ -155,6 +159,7 @@ impl Detector { pixel_convs: (0..3).map(|i| Ok(candle_nn::conv2d(HIDDEN, HIDDEN, 3, conv3, p.pp(format!("conv_layers.{i}")))?)).collect::>()?, pixel_norms: (0..3).map(|i| Ok(candle_nn::group_norm(8, HIDDEN, 1e-5, p.pp(format!("norms.{i}")))?)).collect::>()?, instance_proj: candle_nn::conv2d(HIDDEN, HIDDEN, 1, Conv2dConfig::default(), m.pp("instance_projection"))?, + semantic_proj: candle_nn::conv2d(HIDDEN, 1, 1, Conv2dConfig::default(), m.pp("semantic_projection"))?, mask_embedder: (0..3).map(|i| linear(HIDDEN, HIDDEN, m.pp(format!("mask_embedder.layers.{i}")))).collect::>()?, }) } @@ -268,6 +273,7 @@ impl Detector { prev = norm.forward(&conv.forward(&prev)?)?.relu()?; } let inst = self.instance_proj.forward(&prev)?; // [1, 256, 288, 288] + let semantic = to_vec(&self.semantic_proj.forward(&prev)?)?; let (_, c, mh, mw) = inst.dims4()?; let mut e = queries.clone(); for (i, l) in self.mask_embedder.iter().enumerate() { @@ -285,24 +291,98 @@ impl Detector { presence: to_vec(&presence)?.first().copied().unwrap_or(-10.0), boxes: bv.as_chunks::<4>().0.to_vec(), text, + semantic, }) } } +/// How far below 0 the strongest instance may be where it still takes its share of the semantic mask. +const NEAR: f32 = 4.0; + +/// Gives each of `instances` its share of `semantic` (everything the prompt names, one mask): where an +/// instance is the strongest and near (above −[`NEAR`]), the semantic logits replace its own (it keeps +/// its own where it is surer of itself). An instance's own logits stay close to 0 even well inside it +/// (a person comes out about 90 % selected) and both fade where two instances touch, which leaves a gap +/// between them; the semantic mask is surer inside and runs on across that boundary. Cells far from +/// every instance keep their logits, so things no instance found don't join one. +pub fn share(instances: &mut [Vec], semantic: &[f32]) { + let n = semantic.len(); + if instances.is_empty() || instances.iter().any(|l| l.len() != n) { + return; + } + for (i, s) in semantic.iter().enumerate() { + let Some((k, best)) = instances.iter().enumerate().map(|(k, l)| (k, l[i])).max_by(|a, b| a.1.total_cmp(&b.1)) else { return }; + if best > -NEAR { + let l = &mut instances[k][i]; + *l = if *l > 0.0 { l.max(*s) } else { *s }; + } + } +} + /// The instances of `out` scoring above `threshold` (sigmoid(logit) × sigmoid(presence)), -/// merged into one mask (the per-pixel maximum of their probabilities, as logits). +/// merged into one mask: each its share of the semantic mask ([`share`]), then the per-pixel +/// maximum of their probabilities, as logits. pub fn merge(out: &Output, threshold: f32) -> Option { let s = |x: f32| 1.0 / (1.0 + (-x).exp()); let presence = s(out.presence); let keep: Vec<(usize, f32)> = out.logits.iter().enumerate().map(|(i, l)| (i, s(*l) * presence)).filter(|(_, p)| *p > threshold).collect(); let best = keep.iter().map(|k| k.1).fold(0f32, f32::max); - let mut merged: Option> = None; - for (i, _) in &keep { - let m = out.masks.i(*i).ok().and_then(|t| to_vec(&t).ok())?; - merged = Some(match merged { - None => m, - Some(acc) => acc.iter().zip(&m).map(|(a, b)| a.max(*b)).collect(), - }); - } + let mut masks: Vec> = keep.iter().map(|(i, _)| out.masks.i(*i).ok().and_then(|t| to_vec(&t).ok())).collect::>()?; + share(&mut masks, &out.semantic); + let merged = masks.into_iter().reduce(|acc, m| acc.iter().zip(&m).map(|(a, b)| a.max(*b)).collect()); merged.map(|logits| Prediction { logits, score: best }) } + +#[cfg(test)] +mod tests { + use super::*; + + /// Two people side by side on a row of cells, each only weakly sure of themself and both fading where + /// they touch; the semantic mask is sure of both and runs across. Each takes the semantic logits where + /// it is the strongest instance, a third region nobody is near keeps its own. + #[test] + fn instances_share_the_semantic_mask() { + // cells: a (0..3), the seam (3), b (4..7), floor beside b (7), someone neither found (8), background + let a = vec![2.0, 2.5, 1.0, -0.5, -3.0, -5.0, -6.0, -3.5, -9.0, -9.0]; + let b = vec![-6.0, -5.0, -3.0, -1.0, 1.0, 2.5, 2.0, -3.0, -9.0, -9.0]; + let mut inst = vec![a, b]; + let semantic = [5.0, 6.0, 5.0, 4.0, 5.0, 6.0, 5.0, -7.0, 5.0, -8.0]; + share(&mut inst, &semantic); + assert_eq!(inst[0][..4], [5.0, 6.0, 5.0, 4.0], "a's share, the seam included (a is the stronger there)"); + assert_eq!(inst[1][4..7], [5.0, 6.0, 5.0], "b's share"); + assert_eq!(inst[0][4..7], [-3.0, -5.0, -6.0], "a stays out of b"); + assert_eq!(inst[1][7], -7.0, "the floor near b is surely background"); + assert_eq!([inst[0][8], inst[1][8]], [-9.0, -9.0], "someone neither instance is near joins neither"); + // an instance surer than the semantic mask keeps its own logits + let mut sure = vec![vec![9.0, 1.0]]; + share(&mut sure, &[5.0, 5.0]); + assert_eq!(sure[0], [9.0, 5.0]); + // mismatched sizes (no semantic mask) leave the instances alone + let mut none = vec![vec![1.0, 2.0]]; + share(&mut none, &[]); + assert_eq!(none[0], [1.0, 2.0]); + } + + /// Describe: the instances above the threshold, each with its share of the semantic mask, merged; + /// a query below the threshold adds nothing, and nothing found is `None`. + #[test] + fn describe_merges_the_shares_of_what_it_found() { + // two instances, weak inside and both fading where they touch (cell 2); a third query scores low + let masks = [[2.0f32, 1.0, -0.5, -3.0, -6.0], [-6.0, -3.0, -1.0, 1.0, 2.0], [-9.0, -9.0, -9.0, -9.0, 3.0]]; + let output = |logits: [f32; 3], semantic: &[f32]| Output { + masks: Tensor::from_vec(masks.concat(), (3, 1, 5), &Device::Cpu).unwrap(), + logits: logits.to_vec(), + presence: 8.0, + boxes: vec![[0.0; 4]; 3], + text: Tensor::zeros(1, candle_core::DType::F32, &Device::Cpu).unwrap(), + semantic: semantic.to_vec(), + }; + let semantic = [5.0, 6.0, 4.0, 6.0, 5.0]; + let merged = merge(&output([4.0, 4.0, -6.0], &semantic), 0.5).unwrap(); + assert_eq!(merged.logits, semantic, "sure inside both and across where they touch"); + // without a semantic mask: the per-pixel maximum of the instances, as before + let plain = merge(&output([4.0, 4.0, -6.0], &[]), 0.5).unwrap(); + assert_eq!(plain.logits, [2.0, 1.0, -0.5, 1.0, 2.0]); + assert!(merge(&output([-6.0; 3], &semantic), 0.5).is_none(), "nothing found"); + } +} diff --git a/crates/segment/tests/reference.rs b/crates/segment/tests/reference.rs index 315d36cd1..a2a110053 100644 --- a/crates/segment/tests/reference.rs +++ b/crates/segment/tests/reference.rs @@ -93,5 +93,12 @@ fn matches_reference() { let want = vec(&r[&format!("det{i}_masks")].i((0, q)).unwrap()); assert!(rel_err("best mask", &ours, &want) < 2e-2); assert!(mask_disagreement(&ours, &want) < 2e-3); + // the semantic mask (all instances as one; Subject's people take their share of it); reference + // files made before it was written lack it + if let Some(want) = r.get(&format!("det{i}_semantic")) { + let want = vec(want); + assert!(rel_err("semantic mask", &out.semantic, &want) < 2e-2); + assert!(mask_disagreement(&out.semantic, &want) < 2e-3); + } } } diff --git a/docs/ai-masks.md b/docs/ai-masks.md index b02912abc..f4b3cebef 100644 --- a/docs/ai-masks.md +++ b/docs/ai-masks.md @@ -16,7 +16,10 @@ In the Masking panel (M): to leave it out. Clicks show as green (include) and red (leave out) dots. - **Describe**: type what to select — `sky`, `trees`, `the red car` — and press Return. Every instance the model finds (score above 0.5) joins the mask; if nothing matches, nothing is - created and you are told so. + created and you are told so. Around what it found, the mask follows the model's semantic mask + for the phrase, which is sure inside and runs on where two instances touch (the instances' + own masks alone left a sky or a net about 70–90 % selected and a thin line between two people + side by side). - Several things at once: `car, road` (comma-separated) selects each and merges them. - **+ / −** under the selected mask (and the Add / Subtract / Intersect menus) combine selections: Describe `car`, click the mask, + ▸ Describe… `road`; − ▸ Describe… `people` takes @@ -214,7 +217,7 @@ wait and return the result (when built with the `sam` feature; the default CLI h `NOTICE` — the ported files say so and that they were modified): the 32-layer ViT backbone with 2-D RoPE and windowed attention, the feature pyramids, the SAM 2-style prompt encoder and two-way mask decoder for clicks, and the CLIP text encoder, DETR encoder/decoder (box - relative-position bias, presence token) and pixel decoder for text. The CLIP tokenizer is a + relative-position bias, presence token), pixel decoder and semantic head for text. The CLIP tokenizer is a small BPE in `tokenizer.rs`. The downloader (original code) is the `lightcraft-fetch` crate, shared by any model that is downloaded; `fetch/` here says what SAM 3 needs. - candle is pinned at 0.9.2: later releases make `candle-core` depend on `tokenizers` with the Oniguruma C library, and the product is pure Rust. Weights are read with positional reads (no diff --git a/tools/sam3_reference.py b/tools/sam3_reference.py index 66d9def90..99a06f3ee 100644 --- a/tools/sam3_reference.py +++ b/tools/sam3_reference.py @@ -58,6 +58,7 @@ for i, text in enumerate(["moon", "sky", "bird"]): out[f"det{i}_presence"] = o.presence_logits.contiguous() out[f"det{i}_boxes"] = o.pred_boxes.contiguous() out[f"det{i}_masks"] = o.pred_masks.contiguous() + out[f"det{i}_semantic"] = o.semantic_seg.contiguous() sc = (o.pred_logits.sigmoid() * o.presence_logits.sigmoid())[0] print(text, "presence", o.presence_logits.sigmoid().item(), "top", sc.topk(3).values.tolist(), file=sys.stderr) for i, f in enumerate(det.vision_encoder(px).fpn_hidden_states):