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 <noreply@anthropic.com>

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;

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<f32>)> {
+    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<f32> = (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<Vec<f32>> {
+    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<Tensor> {
+    let side = INPUT_SIZE;
+    let f: Vec<f32> = 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::<f32>().unwrap();
+        assert!(v.iter().all(|x| *x == 0.0 || *x == 1.0));
+        assert!(v.iter().sum::<f32>() > 0.0);
+    }
+}

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<FrameMask> {
+        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::<f32>()?;
+        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<FrameMask> {
+        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::<f32>()?, (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::<f32>()?;
+        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<usize>, cancel: &Cancel, out: &mut dyn FnMut(FrameMask)) -> Result<()> {

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<u8> = 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(),

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::<f32>().unwrap();
+        let want_score = r.get(&format!("maskprompt.{tag}.{}.score", m.frame)).unwrap().flatten_all().unwrap().to_vec1::<f32>().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<std::sync::Arc<ClipFrame>> = (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<f32> = (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());
+}

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():