From 8a6ad6f5e9d5cbc104820919a4712a807cbfe192 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:19:14 +0200 Subject: [PATCH] S2.13: prompt encoder and two-way mask decoder match PyTorch A click on the synthetic image: mask logits within 1.1e-4 (range 11.6), IoU head 1e-6, object score 2e-5, prompt embeddings 1e-7. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/src/decoder.rs | 258 +++++++++++++++++++++++++++++++++ crates/sam2/src/lib.rs | 3 + crates/sam2/src/model.rs | 77 ++++++++++ crates/sam2/src/prompt.rs | 181 +++++++++++++++++++++++ crates/sam2/tests/reference.rs | 36 +++++ 5 files changed, 555 insertions(+) create mode 100644 crates/sam2/src/decoder.rs create mode 100644 crates/sam2/src/model.rs create mode 100644 crates/sam2/src/prompt.rs diff --git a/crates/sam2/src/decoder.rs b/crates/sam2/src/decoder.rs new file mode 100644 index 000000000..e42a7acd9 --- /dev/null +++ b/crates/sam2/src/decoder.rs @@ -0,0 +1,258 @@ +// Modified work: ported from Hugging Face `transformers` `models/sam2/modeling_sam2.py` +// (Sam2Attention, Sam2TwoWayAttentionBlock, Sam2TwoWayTransformer, Sam2MaskDecoder; Apache-2.0, +// Copyright The HuggingFace Inc. team) to Rust on candle. See NOTICE. +//! The two-way mask decoder: output tokens and prompt tokens attend to the image and back, then +//! hyper-networks turn the mask tokens into masks over the up-scaled image embedding. +//! One object per call (batch 1, one point batch). + +use candle_core::{D, Tensor}; +use candle_nn::{LayerNorm, Linear, Module, ops::sigmoid, ops::softmax}; + +use crate::nn::{layer_norm, linear}; +use crate::weights::Weights; +use crate::{Error, Result}; + +const DIM: usize = 256; +const MASK_TOKENS: usize = 4; + +struct Attention { + q: Linear, + k: Linear, + v: Linear, + o: Linear, + heads: usize, + scale: f64, +} + +impl Attention { + fn load(w: &Weights, prefix: &str, downsample: usize, heads: usize) -> Result { + let inner = DIM / downsample; + Ok(Attention { + q: linear(w, &format!("{prefix}.q_proj"), DIM, inner)?, + k: linear(w, &format!("{prefix}.k_proj"), DIM, inner)?, + v: linear(w, &format!("{prefix}.v_proj"), DIM, inner)?, + o: linear(w, &format!("{prefix}.o_proj"), inner, DIM)?, + heads, + scale: ((inner / heads) as f64).powf(-0.5), + }) + } + + /// `q`: `[1, nq, 256]`, `k` and `v`: `[1, nk, 256]`. + fn forward(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result { + let split = |t: Tensor| -> Result { Ok(t.reshape((1, (), self.heads, t.dim(2)? / self.heads))?.transpose(1, 2)?.contiguous()?) }; + let (q, k, v) = (split(self.q.forward(q)?)?, split(self.k.forward(k)?)?, split(self.v.forward(v)?)?); + let att = softmax(&(q.matmul(&k.transpose(2, 3)?.contiguous()?)? * self.scale)?, D::Minus1)?; + let out = att.matmul(&v)?.transpose(1, 2)?.contiguous()?; + let (n, h, d) = (out.dim(1)?, out.dim(2)?, out.dim(3)?); + Ok(self.o.forward(&out.reshape((1, n, h * d))?)?) + } +} + +/// Feed-forward with ReLU, `Sam2FeedForward`: `proj_in`, `layers`, `proj_out`. +pub struct FeedForward { + proj_in: Linear, + layers: Vec, + proj_out: Linear, + sigmoid_output: bool, +} + +impl FeedForward { + pub fn load(w: &Weights, prefix: &str, input: usize, hidden: usize, output: usize, num_layers: usize, sigmoid_output: bool) -> Result { + Ok(FeedForward { + proj_in: linear(w, &format!("{prefix}.proj_in"), input, hidden)?, + layers: (0..num_layers.saturating_sub(2)).map(|i| linear(w, &format!("{prefix}.layers.{i}"), hidden, hidden)).collect::>()?, + proj_out: linear(w, &format!("{prefix}.proj_out"), hidden, output)?, + sigmoid_output, + }) + } + + pub fn forward(&self, x: &Tensor) -> Result { + let mut x = self.proj_in.forward(x)?.relu()?; + for l in &self.layers { + x = l.forward(&x)?.relu()?; + } + let x = self.proj_out.forward(&x)?; + Ok(if self.sigmoid_output { sigmoid(&x)? } else { x }) + } +} + +struct Block { + self_attn: Attention, + ln1: LayerNorm, + token_to_image: Attention, + ln2: LayerNorm, + mlp: FeedForward, + ln3: LayerNorm, + ln4: LayerNorm, + image_to_token: Attention, + skip_first_pe: bool, +} + +impl Block { + fn forward(&self, queries: &Tensor, keys: &Tensor, query_pe: &Tensor, key_pe: &Tensor) -> Result<(Tensor, Tensor)> { + let mut queries = if self.skip_first_pe { + self.self_attn.forward(queries, queries, queries)? + } else { + let q = (queries + query_pe)?; + (queries + self.self_attn.forward(&q, &q, queries)?)? + }; + queries = self.ln1.forward(&queries)?; + let (q, k) = ((&queries + query_pe)?, (keys + key_pe)?); + queries = self.ln2.forward(&(&queries + self.token_to_image.forward(&q, &k, keys)?)?)?; + queries = self.ln3.forward(&(&queries + self.mlp.forward(&queries)?)?)?; + let (q, k) = ((&queries + query_pe)?, (keys + key_pe)?); + let keys = self.ln4.forward(&(keys + self.image_to_token.forward(&k, &q, &queries)?)?)?; + Ok((queries, keys)) + } +} + +/// What the decoder returns for one object. +pub struct DecoderOutput { + /// Mask logits `[1, k, 4 * GRID, 4 * GRID]` (`k` is 3 for multimask output, else 1). + pub masks: Tensor, + /// Predicted IoU per mask `[1, k]`. + pub iou: Tensor, + /// The output tokens of the chosen masks `[1, k, 256]` (the object pointer is made from these). + pub sam_tokens: Tensor, + /// Object-presence logit `[1, 1]` (negative: the object is not in the frame). + pub object_score: Tensor, +} + +pub struct MaskDecoder { + obj_token: Tensor, + iou_token: Tensor, + mask_tokens: Tensor, + layers: Vec, + final_attn: Attention, + final_ln: LayerNorm, + up1: (Tensor, Tensor), + up2: (Tensor, Tensor), + up_ln: LayerNorm, + conv_s0: (Tensor, Tensor), + conv_s1: (Tensor, Tensor), + hyper: Vec, + iou_head: FeedForward, + obj_head: FeedForward, +} + +fn conv_w(w: &Weights, name: &str, i: usize, o: usize, k: usize, transposed: bool) -> Result<(Tensor, Tensor)> { + let shape = if transposed { [i, o, k, k] } else { [o, i, k, k] }; + Ok((w.get(&format!("{name}.weight"), &shape)?, w.get(&format!("{name}.bias"), &[o])?)) +} + +fn add_bias(x: Tensor, b: &Tensor) -> Result { + Ok(x.broadcast_add(&b.reshape((1, b.dim(0)?, 1, 1))?)?) +} + +impl MaskDecoder { + pub fn load(w: &Weights) -> Result { + let layers = (0..2) + .map(|i| { + let p = format!("mask_decoder.transformer.layers.{i}"); + Ok(Block { + self_attn: Attention::load(w, &format!("{p}.self_attn"), 1, 8)?, + ln1: layer_norm(w, &format!("{p}.layer_norm1"), DIM, 1e-5)?, + token_to_image: Attention::load(w, &format!("{p}.cross_attn_token_to_image"), 2, 8)?, + ln2: layer_norm(w, &format!("{p}.layer_norm2"), DIM, 1e-5)?, + mlp: FeedForward::load(w, &format!("{p}.mlp"), DIM, 2048, DIM, 2, false)?, + ln3: layer_norm(w, &format!("{p}.layer_norm3"), DIM, 1e-5)?, + ln4: layer_norm(w, &format!("{p}.layer_norm4"), DIM, 1e-5)?, + image_to_token: Attention::load(w, &format!("{p}.cross_attn_image_to_token"), 2, 8)?, + skip_first_pe: i == 0, + }) + }) + .collect::>>()?; + Ok(MaskDecoder { + obj_token: w.get("mask_decoder.obj_score_token.weight", &[1, DIM])?, + iou_token: w.get("mask_decoder.iou_token.weight", &[1, DIM])?, + mask_tokens: w.get("mask_decoder.mask_tokens.weight", &[MASK_TOKENS, DIM])?, + layers, + final_attn: Attention::load(w, "mask_decoder.transformer.final_attn_token_to_image", 2, 8)?, + final_ln: layer_norm(w, "mask_decoder.transformer.layer_norm_final_attn", DIM, 1e-5)?, + up1: conv_w(w, "mask_decoder.upscale_conv1", DIM, DIM / 4, 2, true)?, + up2: conv_w(w, "mask_decoder.upscale_conv2", DIM / 4, DIM / 8, 2, true)?, + up_ln: layer_norm(w, "mask_decoder.upscale_layer_norm", DIM / 4, 1e-6)?, + conv_s0: conv_w(w, "mask_decoder.conv_s0", DIM, DIM / 8, 1, false)?, + conv_s1: conv_w(w, "mask_decoder.conv_s1", DIM, DIM / 4, 1, false)?, + hyper: (0..MASK_TOKENS) + .map(|i| FeedForward::load(w, &format!("mask_decoder.output_hypernetworks_mlps.{i}"), DIM, DIM, DIM / 8, 3, false)) + .collect::>()?, + iou_head: FeedForward::load(w, "mask_decoder.iou_prediction_head", DIM, DIM, MASK_TOKENS, 3, true)?, + obj_head: FeedForward::load(w, "mask_decoder.pred_obj_score_head", DIM, DIM, 1, 3, false)?, + }) + } + + /// Project the two highest-resolution neck levels to the decoder's skip features + /// (`conv_s0` on the 4x level, `conv_s1` on the 2x level), `[1, 32, 4g, 4g]` and `[1, 64, 2g, 2g]`. + pub fn high_res(&self, level0: &Tensor, level1: &Tensor) -> Result<(Tensor, Tensor)> { + let s0 = add_bias(level0.conv2d(&self.conv_s0.0, 0, 1, 1, 1)?, &self.conv_s0.1)?; + let s1 = add_bias(level1.conv2d(&self.conv_s1.0, 0, 1, 1, 1)?, &self.conv_s1.1)?; + Ok((s0, s1)) + } + + /// Predict masks. `image` and `pos`: `[1, 256, g, g]` (the lowest-resolution level with the + /// no-memory embedding already added when there is no memory), `sparse`: `[1, n, 256]`, + /// `dense`: `[1, 256, g, g]`, `high`: the skip features from [`MaskDecoder::high_res`]. + pub fn forward(&self, image: &Tensor, pos: &Tensor, sparse: &Tensor, dense: &Tensor, high: &(Tensor, Tensor), multimask: bool) -> Result { + let (_, c, g, gw) = image.dims4()?; + if c != DIM || g != gw { + return Err(Error::Input(format!("image embedding must be 256 x g x g, got {:?}", image.dims()))); + } + let output_tokens = Tensor::cat(&[&self.obj_token, &self.iou_token, &self.mask_tokens], 0)?.unsqueeze(0)?; + let tokens = Tensor::cat(&[&output_tokens, sparse], 1)?; + let image = (image + dense)?; + let flat = |t: &Tensor| -> Result { Ok(t.flatten_from(2)?.transpose(1, 2)?.contiguous()?) }; + let (mut keys, key_pe) = (flat(&image)?, flat(pos)?); + let mut queries = tokens.clone(); + for l in &self.layers { + (queries, keys) = l.forward(&queries, &keys, &tokens, &key_pe)?; + } + let (q, k) = ((&queries + &tokens)?, (&keys + &key_pe)?); + let queries = self.final_ln.forward(&(&queries + self.final_attn.forward(&q, &k, &keys)?)?)?; + + let iou_out = queries.narrow(1, 1, 1)?.squeeze(1)?; + let mask_tokens_out = queries.narrow(1, 2, MASK_TOKENS)?; + let obj_out = queries.narrow(1, 0, 1)?.squeeze(1)?; + + let emb = keys.transpose(1, 2)?.contiguous()?.reshape((1, DIM, g, g))?; + let (s0, s1) = high; + let up = add_bias(emb.conv_transpose2d(&self.up1.0, 0, 0, 2, 1)?, &self.up1.1)?; + let up = (up + s1)?; + let up = self.up_ln.forward(&up.permute((0, 2, 3, 1))?)?.permute((0, 3, 1, 2))?.contiguous()?.gelu_erf()?; + let up = (add_bias(up.conv_transpose2d(&self.up2.0, 0, 0, 2, 1)?, &self.up2.1)? + s0)?.gelu_erf()?; + + let hyper: Vec = self.hyper.iter().enumerate().map(|(i, m)| m.forward(&mask_tokens_out.narrow(1, i, 1)?.squeeze(1)?)).collect::>()?; + let hyper = Tensor::stack(&hyper, 1)?; + let (_, uc, uh, uw) = up.dims4()?; + let masks = hyper.matmul(&up.reshape((1, uc, uh * uw))?)?.reshape((1, MASK_TOKENS, uh, uw))?; + let iou = self.iou_head.forward(&iou_out)?; + let object_score = self.obj_head.forward(&obj_out)?; + + let (masks, iou, sam_tokens) = if multimask { + (masks.narrow(1, 1, 3)?, iou.narrow(1, 1, 3)?, mask_tokens_out.narrow(1, 1, 3)?) + } else { + let (m, i) = dynamic_single(&masks, &iou)?; + (m, i, mask_tokens_out.narrow(1, 0, 1)?) + }; + Ok(DecoderOutput { masks, iou, sam_tokens, object_score }) + } +} + +/// Single-mask output with the stability fallback (`_dynamic_multimask_via_stability`): if the +/// first mask token's mask is unstable, use the best of the other three. +fn dynamic_single(masks: &Tensor, iou: &Tensor) -> Result<(Tensor, Tensor)> { + const DELTA: f32 = 0.05; + const THRESH: f32 = 0.98; + let single = masks.narrow(1, 0, 1)?; + let single_iou = iou.narrow(1, 0, 1)?; + let logits = single.flatten_all()?.to_vec1::()?; + let area_i = logits.iter().filter(|v| **v > DELTA).count() as f32; + let area_u = logits.iter().filter(|v| **v > -DELTA).count() as f32; + let stability = if area_u > 0.0 { area_i / area_u } else { 1.0 }; + if stability >= THRESH { + return Ok((single, single_iou)); + } + let multi_iou = iou.narrow(1, 1, 3)?.flatten_all()?.to_vec1::()?; + let best = multi_iou.iter().enumerate().fold((0usize, f32::NEG_INFINITY), |b, (i, v)| if *v > b.1 { (i, *v) } else { b }).0; + Ok((masks.narrow(1, 1 + best, 1)?, iou.narrow(1, 1 + best, 1)?)) +} diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 6515eff7b..0a67f8ad7 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -11,9 +11,12 @@ use std::fmt; +pub mod decoder; pub mod hiera; +pub mod model; pub mod neck; pub mod nn; +pub mod prompt; pub mod weights; /// Everything that can go wrong in this crate. Never a panic. diff --git a/crates/sam2/src/model.rs b/crates/sam2/src/model.rs new file mode 100644 index 000000000..13481fdb9 --- /dev/null +++ b/crates/sam2/src/model.rs @@ -0,0 +1,77 @@ +// Modified work: ported from Hugging Face `transformers` `Sam2Model` (Apache-2.0, Copyright The +// HuggingFace Inc. team) to Rust on candle. See NOTICE. +//! The single-image model: encode a frame once, then decode any number of prompts against it. + +use std::path::Path; + +use candle_core::{Device, Tensor}; + +use crate::decoder::{DecoderOutput, MaskDecoder}; +use crate::neck::{ImageFeatures, VisionEncoder}; +use crate::prompt::{INPUT_SIZE, Prompt, PromptEncoder}; +use crate::weights::Weights; +use crate::{Error, Result}; + +/// A frame after the vision encoder: reusable for every prompt on that frame. +pub struct Encoded { + /// The three neck levels, highest resolution first (`[1, 256, 256, 256]`, `128`, `64`). + pub features: ImageFeatures, + /// The decoder's skip features from the two highest-resolution levels. + pub high: (Tensor, Tensor), +} + +pub struct ImageModel { + pub vision: VisionEncoder, + pub prompt: PromptEncoder, + pub decoder: MaskDecoder, + no_memory: Tensor, + image_pos: Tensor, +} + +impl ImageModel { + pub fn load(w: &Weights) -> Result { + let prompt = PromptEncoder::load(w)?; + Ok(ImageModel { + vision: VisionEncoder::load_tiny(w)?, + decoder: MaskDecoder::load(w)?, + no_memory: w.get("no_memory_embedding", &[1, 1, 256])?.reshape((1, 256, 1, 1))?, + image_pos: prompt.image_positions()?, + prompt, + }) + } + + pub fn open(path: &Path) -> Result { + ImageModel::load(&Weights::open(path)?) + } + + /// Encode a normalised `[1, 3, 1024, 1024]` frame. + pub fn encode(&self, pixels: &Tensor) -> Result { + if pixels.dims() != [1, 3, INPUT_SIZE, INPUT_SIZE] { + return Err(Error::Input(format!("pixels must be [1, 3, {INPUT_SIZE}, {INPUT_SIZE}], got {:?}", pixels.dims()))); + } + let features = self.vision.forward(pixels)?; + let (l0, l1) = match (features.features.first(), features.features.get(1)) { + (Some(a), Some(b)) => (a, b), + _ => return Err(Error::Compute("vision encoder returned too few levels".into())), + }; + let high = self.decoder.high_res(l0, l1)?; + Ok(Encoded { features, high }) + } + + /// Prompt embeddings (sparse tokens, empty dense map) for a click/box prompt. + pub fn embed_prompt(&self, p: &Prompt) -> Result<(Tensor, Tensor)> { + Ok((self.prompt.sparse(p)?, self.prompt.dense_empty()?)) + } + + /// Decode masks for a prompt on an encoded frame. + pub fn segment(&self, enc: &Encoded, p: &Prompt, multimask: bool) -> Result { + let low = enc.features.features.get(2).ok_or_else(|| Error::Compute("vision encoder returned too few levels".into()))?; + let image = low.broadcast_add(&self.no_memory)?; + let (sparse, dense) = self.embed_prompt(p)?; + self.decoder.forward(&image, &self.image_pos, &sparse, &dense, &enc.high, multimask) + } + + pub fn device(&self) -> Device { + Device::Cpu + } +} diff --git a/crates/sam2/src/prompt.rs b/crates/sam2/src/prompt.rs new file mode 100644 index 000000000..95834c3cb --- /dev/null +++ b/crates/sam2/src/prompt.rs @@ -0,0 +1,181 @@ +// Modified work: ported from Hugging Face `transformers` `models/sam2/modeling_sam2.py` +// (Sam2PositionalEmbedding, Sam2MaskEmbedding, Sam2PromptEncoder; Apache-2.0, Copyright The +// HuggingFace Inc. team) to Rust on candle. See NOTICE. +//! The prompt encoder: clicks and boxes become sparse tokens, a mask or "no mask" becomes a dense map. + +use candle_core::{Device, Tensor}; +use candle_nn::{LayerNorm, Module}; + +use crate::nn::layer_norm; +use crate::weights::Weights; +use crate::{Error, Result}; + +/// Side of the (square) model input in pixels; prompts are given in these coordinates. +pub const INPUT_SIZE: usize = 1024; +/// Side of the image-embedding grid (input / 16). +pub const GRID: usize = INPUT_SIZE / 16; +const DIM: usize = 256; + +/// What the user pointed at, in model-input pixels (0..1024). +#[derive(Clone, Debug, Default, PartialEq)] +pub struct Prompt { + /// Clicks and whether each is foreground. + pub points: Vec<([f32; 2], bool)>, + /// A box `[x0, y0, x1, y1]`. + pub bbox: Option<[f32; 4]>, +} + +impl Prompt { + /// Reject non-finite or absurd coordinates (they come from the user or an agent). + pub fn validate(&self) -> Result<()> { + let bad = |v: f32| !v.is_finite() || v.abs() > 4.0 * INPUT_SIZE as f32; + if self.points.iter().any(|(p, _)| p.iter().any(|v| bad(*v))) || self.bbox.is_some_and(|b| b.iter().any(|v| bad(*v))) { + return Err(Error::Input("prompt coordinates must be finite and near the frame".into())); + } + if self.points.len() > 64 { + return Err(Error::Input("at most 64 points per prompt".into())); + } + Ok(()) + } +} + +struct MaskEmbed { + conv: [(Tensor, Tensor, usize, usize); 3], + ln1: LayerNorm, + ln2: LayerNorm, +} + +/// Sparse and dense prompt embeddings. +pub struct PromptEncoder { + gauss: Vec, + point_embed: Vec, + not_a_point: Vec, + no_mask: Tensor, + mask: MaskEmbed, +} + +fn rows(t: Tensor) -> Result> { + Ok(t.flatten_all()?.to_vec1::()?) +} + +impl PromptEncoder { + pub fn load(w: &Weights) -> Result { + let conv = |name: &str, i: usize, o: usize, k: usize, stride: usize| -> Result<(Tensor, Tensor, usize, usize)> { + Ok(( + w.get(&format!("prompt_encoder.mask_embed.{name}.weight"), &[o, i, k, k])?, + w.get(&format!("prompt_encoder.mask_embed.{name}.bias"), &[o])?, + k, + stride, + )) + }; + Ok(PromptEncoder { + gauss: rows(w.get("prompt_encoder.shared_embedding.positional_embedding", &[2, DIM / 2])?)?, + point_embed: rows(w.get("prompt_encoder.point_embed.weight", &[4, DIM])?)?, + not_a_point: rows(w.get("prompt_encoder.not_a_point_embed.weight", &[1, DIM])?)?, + no_mask: w.get("prompt_encoder.no_mask_embed.weight", &[1, DIM])?, + mask: MaskEmbed { + conv: [conv("conv1", 1, 4, 2, 2)?, conv("conv2", 4, 16, 2, 2)?, conv("conv3", 16, DIM, 1, 1)?], + ln1: layer_norm(w, "prompt_encoder.mask_embed.layer_norm1", 4, 1e-6)?, + ln2: layer_norm(w, "prompt_encoder.mask_embed.layer_norm2", 16, 1e-6)?, + }, + }) + } + + /// Random-Fourier position encoding of coordinates already normalised to 0..1: `[n * 256]`. + fn fourier(&self, coords01: &[[f32; 2]]) -> Vec { + let half = DIM / 2; + let mut out = vec![0f32; coords01.len() * DIM]; + for (n, c) in coords01.iter().enumerate() { + let (x, y) = (2.0 * c[0] - 1.0, 2.0 * c[1] - 1.0); + for j in 0..half { + let g0 = self.gauss.get(j).copied().unwrap_or(0.0); + let g1 = self.gauss.get(half + j).copied().unwrap_or(0.0); + let v = 2.0 * std::f32::consts::PI * (x * g0 + y * g1); + if let Some(o) = out.get_mut(n * DIM + j) { + *o = v.sin(); + } + if let Some(o) = out.get_mut(n * DIM + half + j) { + *o = v.cos(); + } + } + } + out + } + + /// The position encoding of the whole `GRID x GRID` embedding, `[1, 256, GRID, GRID]`. + pub fn image_positions(&self) -> Result { + let coords: Vec<[f32; 2]> = (0..GRID * GRID).map(|i| [((i % GRID) as f32 + 0.5) / GRID as f32, ((i / GRID) as f32 + 0.5) / GRID as f32]).collect(); + let flat = self.fourier(&coords); + Ok(Tensor::from_vec(flat, (GRID * GRID, DIM), &Device::Cpu)?.t()?.reshape((1, DIM, GRID, GRID))?.contiguous()?) + } + + fn point_row(&self, i: usize) -> &[f32] { + self.point_embed.get(i * DIM..(i + 1) * DIM).unwrap_or(&[]) + } + + /// Sparse tokens `[1, n, 256]` for a prompt (a click-only prompt gets one padding token). + pub fn sparse(&self, p: &Prompt) -> Result { + p.validate()?; + let norm = |c: [f32; 2]| [(c[0] + 0.5) / INPUT_SIZE as f32, (c[1] + 0.5) / INPUT_SIZE as f32]; + let mut tokens: Vec = vec![]; + let mut n = 0; + let mut push = |v: &[f32]| { + tokens.extend_from_slice(v); + n += 1; + }; + let pts: Vec<[f32; 2]> = p.points.iter().map(|(c, _)| norm(*c)).collect(); + let enc = self.fourier(&pts); + for (k, (_, fg)) in p.points.iter().enumerate() { + let e = enc.get(k * DIM..(k + 1) * DIM).unwrap_or(&[]); + let label_row = self.point_row(usize::from(*fg)); + let v: Vec = e.iter().zip(label_row).map(|(a, b)| a + b).collect(); + push(&v); + } + match p.bbox { + None => push(&self.not_a_point.clone()), + Some(b) => { + let corners = [norm([b[0], b[1]]), norm([b[2], b[3]])]; + let enc = self.fourier(&corners); + for (k, row) in [2usize, 3].into_iter().enumerate() { + let e = enc.get(k * DIM..(k + 1) * DIM).unwrap_or(&[]); + let v: Vec = e.iter().zip(self.point_row(row)).map(|(a, b)| a + b).collect(); + push(&v); + } + push(&self.not_a_point.clone()); + } + } + Ok(Tensor::from_vec(tokens, (1, n, DIM), &Device::Cpu)?) + } + + /// The dense embedding for "no mask prompt", `[1, 256, GRID, GRID]`. + pub fn dense_empty(&self) -> Result { + Ok(self.no_mask.reshape((1, DIM, 1, 1))?.broadcast_as((1, DIM, GRID, GRID))?.contiguous()?) + } + + /// The dense embedding of a mask prompt `[1, 1, 4 * GRID, 4 * GRID]` (low-resolution logits). + pub fn dense_mask(&self, mask: &Tensor) -> Result { + let (c1, c2, c3) = (&self.mask.conv[0], &self.mask.conv[1], &self.mask.conv[2]); + let step = |x: &Tensor, c: &(Tensor, Tensor, usize, usize)| -> Result { + Ok(x.conv2d(&c.0, 0, c.3, 1, 1)?.broadcast_add(&c.1.reshape((1, c.1.dim(0)?, 1, 1))?)?) + }; + let cl = |x: &Tensor, ln: &LayerNorm| -> Result { Ok(ln.forward(&x.permute((0, 2, 3, 1))?)?.permute((0, 3, 1, 2))?.contiguous()?) }; + let x = cl(&step(mask, c1)?, &self.mask.ln1)?.gelu_erf()?; + let x = cl(&step(&x, c2)?, &self.mask.ln2)?.gelu_erf()?; + step(&x, c3) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn hostile_prompts_are_rejected() { + let ok = Prompt { points: vec![([10.0, 20.0], true)], bbox: None }; + assert!(ok.validate().is_ok()); + assert!(Prompt { points: vec![([f32::NAN, 0.0], true)], bbox: None }.validate().is_err()); + assert!(Prompt { points: vec![], bbox: Some([0.0, 0.0, f32::INFINITY, 5.0]) }.validate().is_err()); + assert!(Prompt { points: vec![([1e9, 0.0], false)], bbox: None }.validate().is_err()); + assert!(Prompt { points: vec![([1.0, 1.0], true); 65], bbox: None }.validate().is_err()); + } +} diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index bae7d42e3..6536a9eb5 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -107,3 +107,39 @@ fn neck_and_vision_encoder_match_pytorch() { let e = max_abs(feats.features.first().unwrap(), r.get("vision_encoder.neck.3").unwrap()); assert!(e < NECK_TOL); } + +/// Click-prompted single image: prompt embeddings and decoder outputs against `Sam2Model`. +const DECODER_TOL: f32 = 1e-3; + +#[test] +fn prompt_encoder_and_decoder_match_pytorch() { + use effectcraft_sam2::model::ImageModel; + use effectcraft_sam2::prompt::Prompt; + let Some((w, r)) = load() else { return }; + let model = ImageModel::load(&w).unwrap(); + let pixels = r.get("input.pixel_values").unwrap(); + let prompt = Prompt { points: vec![([460.0, 512.0], true)], bbox: None }; + + let (sparse, dense) = model.embed_prompt(&prompt).unwrap(); + let want_sparse = r.get("prompt_encoder.0").unwrap().squeeze(0).unwrap(); + let e = max_abs(&sparse, &want_sparse); + eprintln!("sparse prompt: max abs error {e:e}"); + assert!(e < 1e-5, "sparse prompt error {e}"); + let e = max_abs(&dense, r.get("prompt_encoder.1").unwrap()); + assert!(e < 1e-6, "dense prompt error {e}"); + + let enc = model.encode(pixels).unwrap(); + let out = model.segment(&enc, &prompt, true).unwrap(); + let want_masks = r.get("mask_decoder.0").unwrap().squeeze(0).unwrap(); + let e = max_abs(&out.masks, &want_masks); + eprintln!("masks: max abs error {e:e} (logit range {:?})", want_masks.flatten_all().unwrap().max(0).unwrap().to_scalar::()); + assert!(e < DECODER_TOL, "mask logits error {e}"); + let e = max_abs(&out.iou, &r.get("mask_decoder.1").unwrap().squeeze(0).unwrap()); + eprintln!("iou: max abs error {e:e}"); + assert!(e < 1e-4, "iou error {e}"); + let e = max_abs(&out.object_score, &r.get("mask_decoder.3").unwrap().squeeze(0).unwrap()); + eprintln!("object score: max abs error {e:e}"); + assert!(e < 1e-4, "object score error {e}"); + let e = max_abs(&out.sam_tokens, &r.get("mask_decoder.2").unwrap().squeeze(0).unwrap()); + assert!(e < 1e-4, "sam token error {e}"); +}