From 81faf5b4b50dcad69b14ee327334f313856eb27d Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 20:34:58 +0200 Subject: [PATCH] S2.66: mask prompts match Sam2VideoModel Antialiased bilinear resize as PyTorch does it, the reference's mask-as-output path (mask is the frame's output, the decoder makes the object pointer), VideoTracker::prompt_mask, and mask-only prompts through the tracker interface. Prompt frame logits exact, tracking forwards and backwards within 0.0019. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/src/lib.rs | 1 + crates/sam2/src/resize.rs | 113 +++++++++++++++++++++++++++++++++ crates/sam2/src/session.rs | 55 ++++++++++++++++ crates/sam2/src/track.rs | 12 ++++ crates/sam2/tests/reference.rs | 88 +++++++++++++++++++++++++ tools/sam2_video_reference.py | 24 +++++++ 6 files changed, 293 insertions(+) create mode 100644 crates/sam2/src/resize.rs diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 6a0eb7fbf..c08ac33b5 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -20,6 +20,7 @@ pub mod model; pub mod neck; pub mod nn; pub mod prompt; +pub mod resize; pub mod session; pub mod track; pub mod weights; diff --git a/crates/sam2/src/resize.rs b/crates/sam2/src/resize.rs new file mode 100644 index 000000000..67e96e9af --- /dev/null +++ b/crates/sam2/src/resize.rs @@ -0,0 +1,113 @@ +// Modified work: the antialiased bilinear resize of PyTorch `F.interpolate(mode="bilinear", +// antialias=True, align_corners=False)`, as used by the Hugging Face `transformers` SAM 2 mask path +// (Apache-2.0 for the model code; the algorithm follows PyTorch's documented behaviour). See NOTICE. +//! Resizing masks to and from the model's resolution. + +use candle_core::{Device, Tensor}; + +use crate::prompt::INPUT_SIZE; +use crate::{Error, Result}; + +/// For each output index: the first input index and the normalised triangle-filter weights. +fn axis_weights(n_in: usize, n_out: usize) -> Vec<(usize, Vec)> { + let scale = n_in as f64 / n_out.max(1) as f64; + let support = if scale >= 1.0 { scale } else { 1.0 }; + let invscale = if scale >= 1.0 { 1.0 / scale } else { 1.0 }; + (0..n_out) + .map(|i| { + let center = scale * (i as f64 + 0.5); + let xmin = ((center - support + 0.5) as i64).max(0) as usize; + let xmax = (((center + support + 0.5) as i64).max(0) as usize).min(n_in); + let mut w: Vec = (xmin..xmax.max(xmin)) + .map(|x| { + let d = ((x as f64 - center + 0.5) * invscale).abs(); + if d < 1.0 { (1.0 - d) as f32 } else { 0.0 } + }) + .collect(); + let total: f32 = w.iter().sum(); + if total > 0.0 { + w.iter_mut().for_each(|v| *v /= total); + } + (xmin, w) + }) + .collect() +} + +/// Antialiased bilinear resize of a row-major `h x w` image to `oh x ow`. +pub fn antialias_bilinear(src: &[f32], (h, w): (usize, usize), (oh, ow): (usize, usize)) -> Result> { + if h == 0 || w == 0 || oh == 0 || ow == 0 || src.len() != h * w { + return Err(Error::Input(format!("cannot resize {} values as {w}x{h} to {ow}x{oh}", src.len()))); + } + let xs = axis_weights(w, ow); + let ys = axis_weights(h, oh); + // horizontal pass: h x ow + let mut mid = vec![0f32; h * ow]; + for y in 0..h { + let row = src.get(y * w..(y + 1) * w).unwrap_or(&[]); + for (x, (x0, wt)) in xs.iter().enumerate() { + let acc: f32 = wt.iter().enumerate().map(|(k, a)| a * row.get(x0 + k).copied().unwrap_or(0.0)).sum(); + if let Some(o) = mid.get_mut(y * ow + x) { + *o = acc; + } + } + } + // vertical pass: oh x ow + let mut out = vec![0f32; oh * ow]; + for (y, (y0, wt)) in ys.iter().enumerate() { + for x in 0..ow { + let acc: f32 = wt.iter().enumerate().map(|(k, a)| a * mid.get((y0 + k) * ow + x).copied().unwrap_or(0.0)).sum(); + if let Some(o) = out.get_mut(y * ow + x) { + *o = acc; + } + } + } + Ok(out) +} + +/// A binary mask of any size as the model's mask prompt: `[1, 1, 1024, 1024]` of 0 and 1 (resized +/// like the reference: antialiased bilinear, then thresholded at one half). +pub fn mask_to_model(mask: &[u8], w: usize, h: usize) -> Result { + let side = INPUT_SIZE; + let f: Vec = mask.iter().map(|v| f32::from(*v != 0)).collect(); + let data = if (w, h) == (side, side) { f } else { antialias_bilinear(&f, (h, w), (side, side))?.into_iter().map(|v| f32::from(v >= 0.5)).collect() }; + Ok(Tensor::from_vec(data, (1, 1, side, side), &Device::Cpu)?) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_constant_stays_constant_and_weights_sum_to_one() { + let out = antialias_bilinear(&[3.0; 64 * 64], (64, 64), (16, 16)).unwrap(); + assert!(out.iter().all(|v| (v - 3.0).abs() < 1e-5)); + let up = antialias_bilinear(&[2.0; 16], (4, 4), (9, 7)).unwrap(); + assert!(up.iter().all(|v| (v - 2.0).abs() < 1e-5)); + } + + #[test] + fn a_half_plane_blurs_over_the_filter_width() { + // 8 x 1: left half 0, right half 1, down to 2 x 1. The triangle filter is four pixels wide on + // each side, so the first output sees the ones at weights 0.375 and 0.125 of a total 3.5. + let src = [0., 0., 0., 0., 1., 1., 1., 1.]; + let out = antialias_bilinear(&src, (1, 8), (1, 2)).unwrap(); + assert!((out[0] - 0.5 / 3.5).abs() < 1e-6 && (out[1] - 3.0 / 3.5).abs() < 1e-6, "{out:?}"); + } + + #[test] + fn hostile_sizes_are_errors() { + assert!(antialias_bilinear(&[1.0; 4], (2, 2), (0, 3)).is_err()); + assert!(antialias_bilinear(&[1.0; 5], (2, 2), (1, 1)).is_err()); + assert!(mask_to_model(&[1; 3], 2, 2).is_err()); + assert!(mask_to_model(&[], 0, 0).is_err()); + } + + #[test] + fn a_small_mask_becomes_a_model_sized_binary_mask() { + let m = mask_to_model(&[0, 1, 1, 0], 2, 2).unwrap(); + assert_eq!(m.dims(), [1, 1, INPUT_SIZE, INPUT_SIZE]); + let v = m.flatten_all().unwrap().to_vec1::().unwrap(); + assert!(v.iter().all(|x| *x == 0.0 || *x == 1.0)); + assert!(v.iter().sum::() > 0.0); + } +} diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index 4b20b5c64..c629a7b20 100644 --- a/crates/sam2/src/session.rs +++ b/crates/sam2/src/session.rs @@ -19,6 +19,7 @@ use crate::memory::{Memory, MemoryEncoder}; use crate::model::ImageModel; use crate::nn::{apply, linear}; use crate::prompt::{INPUT_SIZE, Prompt}; +use crate::resize::antialias_bilinear; use crate::weights::Weights; use crate::{Error, Result}; @@ -87,6 +88,8 @@ pub struct VideoModel { no_object_pointer: Tensor, pointer_proj: FeedForward, pointer_time_proj: Linear, + /// `mask_downsample`: a 4 x 4 stride-4 convolution taking a mask prompt to the prompt encoder's resolution. + mask_down: (Tensor, Tensor), } impl VideoModel { @@ -99,6 +102,7 @@ impl VideoModel { no_object_pointer: w.get("no_object_pointer", &[1, 256])?, pointer_proj: FeedForward::load(w, "object_pointer_proj", 256, 256, 256, 3, false)?, pointer_time_proj: linear(w, "temporal_positional_encoding_projection_layer", 256, 64)?, + mask_down: (w.get("mask_downsample.weight", &[1, 1, 4, 4])?, w.get("mask_downsample.bias", &[1])?), }) } @@ -331,6 +335,57 @@ impl<'a> VideoTracker<'a> { self.step(object, frame, Some(prompt), false, true) } + /// Prompt an object on a frame with a mask: `[1, 1, 1024, 1024]` of 0 and 1 (see + /// [`crate::resize::mask_to_model`]). The mask becomes the frame's output as it is (no decode of + /// the mask itself); the decoder only makes the object pointer. Like a click prompt, the frame is a + /// conditioning frame unless it has been tracked already. + pub fn prompt_mask(&mut self, frame: usize, object: ObjectId, mask: &Tensor) -> Result { + self.check_frame(frame)?; + if mask.dims() != [1, 1, INPUT_SIZE, INPUT_SIZE] { + return Err(Error::Input(format!("a mask prompt must be [1, 1, {INPUT_SIZE}, {INPUT_SIZE}], got {:?}", mask.dims()))); + } + let values = mask.flatten_all()?.to_vec1::()?; + if values.iter().any(|v| !v.is_finite()) { + return Err(Error::Input("a mask prompt must be finite".into())); + } + let tracked = self.objects.get(&object).is_some_and(|o| o.non_cond.contains_key(&frame)); + self.objects.entry(object).or_default(); + self.step_mask(object, frame, mask, &values, !tracked) + } + + /// The reference's `_use_mask_as_output`: the mask is the output, the decoder makes the pointer. + fn step_mask(&mut self, id: ObjectId, frame: usize, mask: &Tensor, values: &[f32], conditioning: bool) -> Result { + let feats = self.features(frame)?; + let model = self.model.clone(); + let present = values.iter().any(|v| *v > 0.0); + let high = ((mask * 20.0)? - 10.0)?; + let low = antialias_bilinear(&high.flatten_all()?.to_vec1::()?, (INPUT_SIZE, INPUT_SIZE), (4 * GRID, 4 * GRID))?; + // The pointer: the decoder with the mask as its dense prompt and no clicks, on the frame's features + // without the no-memory embedding. + let down = mask.conv2d(&model.mask_down.0, 0, 4, 1, 1)?.broadcast_add(&model.mask_down.1.reshape((1, 1, 1, 1))?)?; + let dense = model.image.prompt.dense_mask(&down)?; + let out = model.image.decode_embedded(&feats.low, &feats.high, &model.image.prompt.sparse_none()?, &dense, true)?; + let ious = out.iou.flatten_all()?.to_vec1::()?; + let best = ious.iter().enumerate().fold((0usize, f32::NEG_INFINITY), |b, (i, v)| if *v > b.1 { (i, *v) } else { b }).0; + let token = out.sam_tokens.narrow(1, best, 1)?.squeeze(1)?; + let decoder_present = scalar(&out.object_score)? > 0.0; + let pointer = if decoder_present { model.pointer_proj.forward(&token)? } else { model.no_object_pointer.clone() }; + let pointer = if present { pointer } else { model.no_object_pointer.clone() }; + let memory = model.memory_encoder.memory(&feats.low, &high, present)?; + let score = if present { 10.0 } else { -10.0 }; + let result = FrameMask { frame, object: id, logits: low, iou: 1.0, object_score: score, occluded: !present }; + let state = self.objects.entry(id).or_default(); + let entry = FrameOut { memory, pointer, result: conditioning.then(|| result.clone()) }; + if conditioning { + state.non_cond.remove(&frame); + state.cond.insert(frame, entry); + } else { + state.non_cond.insert(frame, entry); + Self::prune(state, frame); + } + Ok(result) + } + /// Track every object from `start` through `max_frames` frames (the rest of the clip when /// `None`) in `direction`, calling `out` for each object on each frame as it is finished. pub fn propagate(&mut self, start: usize, direction: Direction, max_frames: Option, cancel: &Cancel, out: &mut dyn FnMut(FrameMask)) -> Result<()> { diff --git a/crates/sam2/src/track.rs b/crates/sam2/src/track.rs index a46f7f6c6..2b06d08cd 100644 --- a/crates/sam2/src/track.rs +++ b/crates/sam2/src/track.rs @@ -13,6 +13,7 @@ use effectcraft_segment::{ClipFrame, ObjectTracker, TrackRequest, TrackedFrame}; use crate::mask::{Sam2Mask, preprocess}; use crate::prompt::{INPUT_SIZE, Prompt}; +use crate::resize::mask_to_model; use crate::session::{Cancel, Direction, FrameMask, FrameSource, ObjectId, VideoTracker}; use crate::{Error, Result}; @@ -89,6 +90,17 @@ impl Sam2Mask { return Err(Error::Cancelled); } let (w, h) = sizes.borrow().get(&idx).copied().ok_or_else(|| Error::Compute("frame size unknown".into()))?; + // A prompt that is only a mask (probability per pixel of the frame) prompts with the mask. + if let (true, true, Some(m)) = (p.points.is_empty(), p.bbox.is_none(), p.mask.as_ref()) { + if m.len() != w * h { + return Err(Error::Input(format!("a mask prompt of {} values does not fit a {w}x{h} frame", m.len()))); + } + let binary: Vec = m.iter().map(|v| u8::from(*v >= 0.5)).collect(); + tracker.prompt_mask(idx, object_id(oi), &mask_to_model(&binary, w, h)?)?; + let e = first_of.entry(oi).or_insert(idx); + *e = (*e).min(idx); + continue; + } let (sx, sy) = (INPUT_SIZE as f32 / w as f32, INPUT_SIZE as f32 / h as f32); let prompt = Prompt { points: p.points.iter().map(|(c, fg)| ([c[0] * sx, c[1] * sy], *fg)).collect(), diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index da7dd4706..4d4377374 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -456,3 +456,91 @@ fn sam2_tracks_two_objects_in_one_pass() { let empty = TrackRequest { first: 10, count: n, objects: vec![] }; assert!(model.tracker().unwrap().track(&empty, &mut get, &|| false, &mut |_, _, _| {}).is_err()); } + +/// A mask prompt instead of clicks: the mask is the prompt frame's output, and tracking forwards and +/// backwards from it matches `Sam2VideoModel` given the same mask. +#[test] +fn mask_prompted_tracking_matches_sam2_video_model() { + use effectcraft_sam2::session::{Cancel, Direction, ObjectId, VideoModel, VideoTracker}; + let Some((w, r)) = load_video() else { return }; + let model = std::sync::Arc::new(VideoModel::load(&w).unwrap()); + let mut tracker = VideoTracker::new(model, Box::new(Clip(r.get("clip.pixel_values").unwrap().clone()))).unwrap(); + let obj = ObjectId(1); + let first = tracker.prompt_mask(2, obj, r.get("maskprompt.input").unwrap()).unwrap(); + let check = |tag: &str, m: &effectcraft_sam2::session::FrameMask| { + let want = r.get(&format!("maskprompt.{tag}.{}.pred_masks", m.frame)).unwrap().flatten_all().unwrap().to_vec1::().unwrap(); + let want_score = r.get(&format!("maskprompt.{tag}.{}.score", m.frame)).unwrap().flatten_all().unwrap().to_vec1::().unwrap()[0]; + let err = m.logits.iter().zip(&want).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max); + let j = iou(&m.logits, &want); + eprintln!("mask prompt {tag} frame {}: logit error {err:.4}, IoU {j:.4}, score {:.3} (want {want_score:.3})", m.frame, m.object_score); + assert!(j >= VIDEO_MIN_IOU, "{tag} frame {} IoU {j}", m.frame); + assert!(err < VIDEO_LOGIT_TOL, "{tag} frame {} logit error {err}", m.frame); + assert!((m.object_score - want_score).abs() < 0.05, "{tag} frame {} score {} vs {want_score}", m.frame, m.object_score); + }; + check("prompt", &first); + let cancel = Cancel::new(); + let mut fwd = vec![]; + tracker.propagate(2, Direction::Forward, None, &cancel, &mut |m| fwd.push(m)).unwrap(); + fwd.iter().for_each(|m| check("fwd", m)); + let mut rev = vec![]; + tracker.propagate(2, Direction::Backward, None, &cancel, &mut |m| rev.push(m)).unwrap(); + rev.iter().for_each(|m| check("rev", m)); + assert_eq!((fwd.len(), rev.len()), (5, 3)); + // hostile masks are errors + let bad = Tensor::zeros((1, 1, 8, 8), candle_core::DType::F32, &Device::Cpu).unwrap(); + assert!(tracker.prompt_mask(2, obj, &bad).is_err()); + let nan = (Tensor::zeros((1, 1, 1024, 1024), candle_core::DType::F32, &Device::Cpu).unwrap() + f64::NAN).unwrap(); + assert!(tracker.prompt_mask(2, obj, &nan).is_err()); + assert!(tracker.prompt_mask(99, obj, r.get("maskprompt.input").unwrap()).is_err()); +} + +/// The tracker interface takes a mask as a prompt: the disc's true mask on frame 1 tracks the disc. +#[test] +fn a_mask_prompt_through_the_object_tracker_interface() { + use effectcraft_segment::{ClipFrame, MaskModel, Prompt, TrackRequest}; + let Some((w, _)) = load() else { return }; + let (fw, fh, n) = (128usize, 96usize, 4usize); + let centre = |t: usize| (40.0 + 8.0 * t as f32, 48.0f32); + let inside = |t: usize, x: usize, y: usize| ((x as f32 - centre(t).0).powi(2) + (y as f32 - centre(t).1).powi(2)).sqrt() < 26.0; + let frames: Vec> = (0..n) + .map(|t| { + let mut rgb = vec![[0f32; 3]; fw * fh]; + for y in 0..fh { + for x in 0..fw { + let bg = 0.3 + 0.2 * (x as f32 / 9.0).sin() * (y as f32 / 7.0).cos(); + rgb[y * fw + x] = if inside(t, x, y) { [0.85, 0.4, 0.15] } else { [bg, bg * 0.8 + 0.1, bg * 0.6 + 0.3] }; + } + } + std::sync::Arc::new(ClipFrame { rgb, w: fw, h: fh }) + }) + .collect(); + let model = effectcraft_sam2::mask::Sam2Mask::new(std::sync::Arc::new(effectcraft_sam2::session::VideoModel::load(&w).unwrap())); + let mask: Vec = (0..fw * fh).map(|i| f32::from(inside(1, i % fw, i / fw))).collect(); + let mut get = |f: i64| frames.get(f as usize).cloned(); + let req = TrackRequest::single(0, n, vec![(1, Prompt { points: vec![], bbox: None, mask: Some(mask) })]); + let mut seen = std::collections::BTreeMap::new(); + model + .tracker() + .unwrap() + .track(&req, &mut get, &|| false, &mut |_, f, t| { + seen.insert(f, t); + }) + .unwrap(); + assert_eq!(seen.len(), n); + for (f, t) in &seen { + let (mut i, mut u) = (0usize, 0usize); + for y in 0..fh { + for x in 0..fw { + let (a, b) = (t.prob[y * fw + x] > 0.5, inside(*f as usize, x, y)); + i += usize::from(a && b); + u += usize::from(a || b); + } + } + let j = i as f64 / u.max(1) as f64; + eprintln!("mask-prompted frame {f}: IoU {j:.3}"); + assert!(j > 0.85 && !t.occluded, "frame {f}: IoU {j}"); + } + // a mask that does not fit the frame is an error + let bad = TrackRequest::single(0, n, vec![(1, Prompt { points: vec![], bbox: None, mask: Some(vec![1.0; 5]) })]); + assert!(model.tracker().unwrap().track(&bad, &mut get, &|| false, &mut |_, _, _| {}).is_err()); +} diff --git a/tools/sam2_video_reference.py b/tools/sam2_video_reference.py index 6e769d955..a05223be9 100644 --- a/tools/sam2_video_reference.py +++ b/tools/sam2_video_reference.py @@ -96,6 +96,30 @@ def main(model_dir, out): for o in model.propagate_in_video_iterator(sess, start_frame_idx=PROMPT_FRAME, reverse=True): record("rev", o) + # The same object prompted with a mask instead of a click (binary, at the model's 1024 x 1024). + msess = Sam2VideoInferenceSession(video=pixels, video_height=SIZE, video_width=SIZE, dtype=torch.float32) + mobj = msess.obj_id_to_idx(1) + yy, xx = np.mgrid[0:1024, 0:1024].astype(np.float32) + cx, cy = (80 + 12 * PROMPT_FRAME) * scale, 128 * scale + disc = torch.from_numpy((np.hypot(xx + 0.5 - cx, yy + 0.5 - cy) < 34 * scale).astype(np.float32))[None, None] + taps["maskprompt.input"] = disc.clone() + msess.add_mask_inputs(mobj, PROMPT_FRAME, disc) + msess.obj_with_new_inputs = [1] + with torch.no_grad(): + o = model(msess, frame_idx=PROMPT_FRAME) + taps[f"maskprompt.prompt.{o.frame_idx}.pred_masks"] = o.pred_masks.detach().float().clone() + taps[f"maskprompt.prompt.{o.frame_idx}.score"] = o.object_score_logits.detach().float().clone() + for o in model.propagate_in_video_iterator(msess, start_frame_idx=PROMPT_FRAME): + taps[f"maskprompt.fwd.{o.frame_idx}.pred_masks"] = o.pred_masks.detach().float().clone() + taps[f"maskprompt.fwd.{o.frame_idx}.score"] = o.object_score_logits.detach().float().clone() + for o in model.propagate_in_video_iterator(msess, start_frame_idx=PROMPT_FRAME, reverse=True): + taps[f"maskprompt.rev.{o.frame_idx}.pred_masks"] = o.pred_masks.detach().float().clone() + taps[f"maskprompt.rev.{o.frame_idx}.score"] = o.object_score_logits.detach().float().clone() + md = msess.output_dict_per_obj[mobj] + for kind in ("cond_frame_outputs", "non_cond_frame_outputs"): + for f, oo in md[kind].items(): + taps[f"maskstate.{f}.object_pointer"] = oo["object_pointer"].detach().float().clone() + d = sess.output_dict_per_obj[obj_idx] for kind in ("cond_frame_outputs", "non_cond_frame_outputs"): for f, o in d[kind].items():