Skip to main content

cortiq_engine/
sam3vid.rs

1//! SAM3 **video** mode: the image detector plus the SAM2-style memory-bank
2//! tracker packed in the same CMF (`trk.*`, 309 tensors, plus the tracker's
3//! own FPN neck, which lives in the `fpn.sam2_convs` section).
4//!
5//! One frame goes through the same six steps as transformers 5.20.2
6//! `Sam3VideoModel._det_track_one_frame`:
7//!
8//! 1. **detect** — one image forward ([`Sam3::forward`](crate::sam3::Sam3)),
9//!    score `sigmoid(logit)·sigmoid(presence)`, mask-IoU NMS at
10//!    `det_nms_thresh`, keep above `score_threshold_detection`;
11//! 2. **tracker features** — trunk map → `fpn.sam2_convs` neck → levels
12//!    288/144/72 with `conv_s0`/`conv_s1` on the two finest;
13//! 3. **propagate** — per tracked object the current frame's features are
14//!    fused with the memory bank (`trk.transformer.encoder.*`: 4 RoPE
15//!    layers over past `maskmem_features` + object pointers) and decoded by
16//!    the SAM mask decoder into 288² mask logits, object score and object
17//!    pointer; small holes/sprinkles are filled;
18//! 4. **plan** — IoU association with the detector
19//!    (`assoc_iou_thresh` / `trk_assoc_iou_thresh` / `new_det_thresh`),
20//!    hotstart keep-alive + duplicate bookkeeping, recent-occlusion
21//!    suppression, non-overlap constraints and a fresh memory encode of the
22//!    resulting masks;
23//! 5. **execute** — new detections become conditioning frames
24//!    (`_use_mask_as_output`: the detection mask IS the output, the decoder
25//!    only provides the object pointer), dropped objects leave the session;
26//! 6. **outputs** — per-object 288² mask logits, bilinearly interpolated to
27//!    the frame and binarized at 0; box = mask bounding box (inclusive
28//!    corners, like `torchvision.masks_to_boxes`); `track_id` is the session
29//!    object id, stable for the object's whole life.
30//!
31//! Batch size 1 — one object at a time, as in the reference.
32
33use std::collections::{BTreeMap, HashMap, HashSet};
34use std::time::Instant;
35
36use anyhow::{anyhow, Result};
37use cortiq_core::CmfModel;
38
39use crate::inference::gelu_erf;
40use crate::pool::Pool;
41use crate::sam3::{self, Sam3, Sam3Out};
42use crate::sam3vid_net::{
43    dense, ln_chw, resize2d, sine, sine_1d, sigm, Net, TrkFeats, EMB, HID, MEM, NO_OBJ_SCORE,
44};
45
46/// Low-resolution mask side (`low_res_mask_size`), also the detector's
47/// `cfg.processor.mask_size`.
48pub const MASK_SIDE: usize = 288;
49/// Tracker input resolution (`tracker_config.image_size`).
50const IMG: usize = 1008;
51/// Memory-encoder mask target side (`mask_input_size · 4`).
52const MEM_SIDE: usize = 1152;
53
54// ───────────────────── configuration ─────────────────────
55
56/// Every knob the video pipeline reads, from the CMF's `cfg.tracker`
57/// (the vendor `config.json` of `Sam3VideoModel`). The defaults below are
58/// the shipped values; [`VidCfg::from_cmf`] overrides them from the file.
59#[derive(Clone, Debug, PartialEq)]
60pub struct VidCfg {
61    pub score_threshold_detection: f32,
62    pub det_nms_thresh: f32,
63    pub assoc_iou_thresh: f32,
64    pub trk_assoc_iou_thresh: f32,
65    pub new_det_thresh: f32,
66    pub recondition_on_trk_masks: bool,
67    pub recondition_every_nth_frame: i64,
68    pub hotstart_delay: i64,
69    pub hotstart_unmatch_thresh: usize,
70    pub hotstart_dup_thresh: usize,
71    pub suppress_unmatched_only_within_hotstart: bool,
72    pub init_trk_keep_alive: i64,
73    pub max_trk_keep_alive: i64,
74    pub min_trk_keep_alive: i64,
75    pub suppress_recent_occlusion_thresh: f32,
76    pub decrease_keep_alive_for_empty: bool,
77    pub fill_hole_area: usize,
78    pub max_num_objects: usize,
79    pub high_conf_thresh: f32,
80    pub high_iou_thresh: f32,
81    pub low_res_mask_size: usize,
82    pub num_maskmem: usize,
83    pub max_cond_frame_num: usize,
84    pub max_object_pointers: usize,
85    pub sigmoid_scale: f32,
86    pub sigmoid_bias: f32,
87    pub temporal_ptr_pe: bool,
88    pub occlusion_spatial: bool,
89    /// `multimask_output_in_sam && multimask_output_for_tracking` — the
90    /// tracker decodes three masks and keeps the best-IoU one.
91    pub multimask: bool,
92}
93
94impl Default for VidCfg {
95    fn default() -> Self {
96        Self {
97            score_threshold_detection: 0.5,
98            det_nms_thresh: 0.1,
99            assoc_iou_thresh: 0.1,
100            trk_assoc_iou_thresh: 0.5,
101            new_det_thresh: 0.7,
102            recondition_on_trk_masks: false,
103            recondition_every_nth_frame: 16,
104            hotstart_delay: 15,
105            hotstart_unmatch_thresh: 8,
106            hotstart_dup_thresh: 8,
107            suppress_unmatched_only_within_hotstart: true,
108            init_trk_keep_alive: 30,
109            max_trk_keep_alive: 30,
110            min_trk_keep_alive: -1,
111            suppress_recent_occlusion_thresh: 0.7,
112            decrease_keep_alive_for_empty: false,
113            fill_hole_area: 16,
114            max_num_objects: 10_000,
115            high_conf_thresh: 0.8,
116            high_iou_thresh: 0.8,
117            low_res_mask_size: MASK_SIDE,
118            num_maskmem: 7,
119            max_cond_frame_num: 4,
120            max_object_pointers: 16,
121            sigmoid_scale: 20.0,
122            sigmoid_bias: -10.0,
123            temporal_ptr_pe: true,
124            occlusion_spatial: true,
125            multimask: true,
126        }
127    }
128}
129
130impl VidCfg {
131    /// Parse `cfg.tracker` out of the CMF (missing keys keep the defaults).
132    pub fn from_cmf(m: &CmfModel) -> Result<Self> {
133        let mut c = Self::default();
134        let raw = m
135            .tensor("cfg.tracker")
136            .ok_or_else(|| anyhow!("the CMF carries no cfg.tracker section — not a SAM3 video model"))?;
137        let bytes = m.entry_bytes(raw);
138        let v: serde_json::Value =
139            serde_json::from_slice(bytes).map_err(|e| anyhow!("cfg.tracker is not JSON: {e}"))?;
140        let f = |k: &str, d: f32| v[k].as_f64().map_or(d, |x| x as f32);
141        let i = |k: &str, d: i64| v[k].as_i64().unwrap_or(d);
142        let b = |k: &str, d: bool| v[k].as_bool().unwrap_or(d);
143        let u = |k: &str, d: usize| v[k].as_u64().map_or(d, |x| x as usize);
144
145        c.score_threshold_detection = f("score_threshold_detection", c.score_threshold_detection);
146        c.det_nms_thresh = f("det_nms_thresh", c.det_nms_thresh);
147        c.assoc_iou_thresh = f("assoc_iou_thresh", c.assoc_iou_thresh);
148        c.trk_assoc_iou_thresh = f("trk_assoc_iou_thresh", c.trk_assoc_iou_thresh);
149        c.new_det_thresh = f("new_det_thresh", c.new_det_thresh);
150        c.recondition_on_trk_masks = b("recondition_on_trk_masks", c.recondition_on_trk_masks);
151        c.recondition_every_nth_frame = i("recondition_every_nth_frame", c.recondition_every_nth_frame);
152        c.hotstart_delay = i("hotstart_delay", c.hotstart_delay);
153        c.hotstart_unmatch_thresh = u("hotstart_unmatch_thresh", c.hotstart_unmatch_thresh);
154        c.hotstart_dup_thresh = u("hotstart_dup_thresh", c.hotstart_dup_thresh);
155        c.suppress_unmatched_only_within_hotstart =
156            b("suppress_unmatched_only_within_hotstart", c.suppress_unmatched_only_within_hotstart);
157        c.init_trk_keep_alive = i("init_trk_keep_alive", c.init_trk_keep_alive);
158        c.max_trk_keep_alive = i("max_trk_keep_alive", c.max_trk_keep_alive);
159        c.min_trk_keep_alive = i("min_trk_keep_alive", c.min_trk_keep_alive);
160        c.suppress_recent_occlusion_thresh = f(
161            "suppress_overlapping_based_on_recent_occlusion_threshold",
162            c.suppress_recent_occlusion_thresh,
163        );
164        c.decrease_keep_alive_for_empty =
165            b("decrease_trk_keep_alive_for_empty_masklets", c.decrease_keep_alive_for_empty);
166        c.fill_hole_area = u("fill_hole_area", c.fill_hole_area);
167        c.max_num_objects = u("max_num_objects", c.max_num_objects);
168        c.high_conf_thresh = f("high_conf_thresh", c.high_conf_thresh);
169        c.high_iou_thresh = f("high_iou_thresh", c.high_iou_thresh);
170        c.low_res_mask_size = u("low_res_mask_size", c.low_res_mask_size);
171        let t = &v["tracker_config"];
172        c.num_maskmem = t["num_maskmem"].as_u64().map_or(c.num_maskmem, |x| x as usize);
173        c.max_cond_frame_num = t["max_cond_frame_num"].as_u64().map_or(c.max_cond_frame_num, |x| x as usize);
174        c.max_object_pointers =
175            t["max_object_pointers_in_encoder"].as_u64().map_or(c.max_object_pointers, |x| x as usize);
176        c.sigmoid_scale = t["sigmoid_scale_for_mem_enc"].as_f64().map_or(c.sigmoid_scale, |x| x as f32);
177        c.sigmoid_bias = t["sigmoid_bias_for_mem_enc"].as_f64().map_or(c.sigmoid_bias, |x| x as f32);
178        c.temporal_ptr_pe = t["enable_temporal_pos_encoding_for_object_pointers"]
179            .as_bool()
180            .unwrap_or(c.temporal_ptr_pe);
181        c.occlusion_spatial =
182            t["enable_occlusion_spatial_embedding"].as_bool().unwrap_or(c.occlusion_spatial);
183        c.multimask = t["multimask_output_in_sam"].as_bool().unwrap_or(true)
184            && t["multimask_output_for_tracking"].as_bool().unwrap_or(true);
185        anyhow::ensure!(
186            c.low_res_mask_size == MASK_SIDE,
187            "unsupported low_res_mask_size {}",
188            c.low_res_mask_size
189        );
190        anyhow::ensure!(c.num_maskmem <= 7, "num_maskmem {} > 7", c.num_maskmem);
191        Ok(c)
192    }
193}
194
195// ───────────────────── session state ─────────────────────
196
197/// One object's stored output for one frame.
198#[derive(Default, Clone)]
199struct ObjOut {
200    /// 288² mask logits.
201    pred_masks: Vec<f32>,
202    /// Decoder object-score logit (±10 for mask-input frames).
203    obj_score: f32,
204    /// 256-d object pointer.
205    pointer: Vec<f32>,
206    mem_feat: Option<Vec<f32>>,
207    mem_pos: Option<Vec<f32>>,
208}
209
210#[derive(Default)]
211struct ObjState {
212    cond: BTreeMap<usize, ObjOut>,
213    non_cond: BTreeMap<usize, ObjOut>,
214}
215
216/// Persistent tracking state across frames (`Sam3VideoInferenceSession`).
217pub struct Session {
218    /// Frames in the clip (0 = unknown/streaming); the reference caps
219    /// `max_object_pointers_to_use` with it.
220    pub num_frames: usize,
221    /// Ordered object ids — the **track ids** the CLI reports.
222    pub obj_ids: Vec<i64>,
223    pub max_obj_id: i64,
224    objs: HashMap<i64, ObjState>,
225    /// Detection score captured when the object was created.
226    score: HashMap<i64, f32>,
227    /// Frame-wise tracker score (sigmoid of the object-score logit).
228    trk_score_fw: HashMap<i64, HashMap<usize, f32>>,
229    last_occluded: HashMap<i64, i64>,
230    first_frame: HashMap<i64, i64>,
231    unmatched: HashMap<i64, Vec<usize>>,
232    keep_alive: HashMap<i64, i64>,
233    overlap: HashMap<(i64, i64), Vec<usize>>,
234    removed: HashSet<i64>,
235    suppressed: HashSet<i64>,
236}
237
238impl Default for Session {
239    fn default() -> Self {
240        Self {
241            // the reference starts `max_obj_id` at -1, so the first object
242            // created gets track id 0
243            max_obj_id: -1,
244            num_frames: 0,
245            obj_ids: Vec::new(),
246            objs: HashMap::new(),
247            score: HashMap::new(),
248            trk_score_fw: HashMap::new(),
249            last_occluded: HashMap::new(),
250            first_frame: HashMap::new(),
251            unmatched: HashMap::new(),
252            keep_alive: HashMap::new(),
253            overlap: HashMap::new(),
254            removed: HashSet::new(),
255            suppressed: HashSet::new(),
256        }
257    }
258}
259
260impl Session {
261    pub fn new() -> Self {
262        Self::default()
263    }
264
265    /// Object ids hidden from the outputs (suppressed / hotstart-removed).
266    pub fn hidden_ids(&self) -> Vec<i64> {
267        self.removed.union(&self.suppressed).copied().collect()
268    }
269}
270
271// ───────────────────── outputs ─────────────────────
272
273/// One object's observation on one frame.
274#[derive(Clone, Debug)]
275pub struct TrackObs {
276    pub id: i64,
277    pub score: f32,
278    /// `xyxy` in frame pixels (inclusive corners), from the mask box.
279    pub box_px: [f32; 4],
280    /// Binary mask at frame resolution, row-major `width × height`.
281    pub mask: Vec<u8>,
282    /// `true` when this frame's mask came straight from the detector.
283    pub from_detection: bool,
284    /// `true` when the tracker propagated the object onto this frame.
285    pub from_tracker: bool,
286}
287
288/// Everything one frame produced.
289#[derive(Clone, Debug)]
290pub struct FrameOut {
291    pub frame: usize,
292    pub tracks: Vec<TrackObs>,
293    /// Raw detections above threshold this frame (before association).
294    pub detections: usize,
295    pub presence: f32,
296    pub det_ms: f64,
297    pub trk_ms: f64,
298    pub total_ms: f64,
299}
300
301// ───────────────────── the model ─────────────────────
302
303struct Det {
304    score: f32,
305    mask: Vec<f32>,
306}
307
308/// Detector + tracker over one CMF.
309pub struct Sam3Video {
310    sam: Sam3,
311    pub cfg: VidCfg,
312}
313
314impl Sam3Video {
315    pub fn open(path: &str) -> Result<Self> {
316        let sam = Sam3::open(path)?;
317        let cfg = VidCfg::from_cmf(sam.cmf())?;
318        Ok(Self { sam, cfg })
319    }
320
321    /// The image detector (shared prompt text tower, timing probes).
322    pub fn detector(&mut self) -> &mut Sam3 {
323        &mut self.sam
324    }
325
326    fn net(&self) -> Net<'_> {
327        Net {
328            m: self.sam.cmf(),
329            pool: self.sam.pool_handle().map(|p| p.as_ref()),
330            cfg: &self.cfg,
331        }
332    }
333
334    /// Run one frame: detector → propagation → planning → new objects →
335    /// outputs. `frame_size` is the original `(width, height)`.
336    pub fn forward(
337        &mut self,
338        s: &mut Session,
339        frame_idx: usize,
340        pixels: &[f32],
341        ids: &[u32],
342        attn: &[u8],
343        frame_size: (usize, usize),
344    ) -> Result<FrameOut> {
345        let _prof = crate::cpuprof::spans::dump_on_drop(format!("video frame {frame_idx}"));
346        let t0 = Instant::now();
347        // ── 1. detector ────────────────────────────────────────────
348        let t_fwd = Instant::now();
349        let out = self.sam.forward(pixels, ids, attn)?;
350        let mut det_ms = t_fwd.elapsed().as_secs_f64() * 1e3;
351        let t_det = Instant::now();
352        let _pd = crate::cpuprof::spans::span("vid.detect");
353        let dets = self.detect(&out)?;
354        drop(_pd);
355        det_ms += t_det.elapsed().as_secs_f64() * 1e3;
356        let presence = sigm(out.presence_logit);
357
358        let net = self.net();
359        let t_trk = Instant::now();
360
361        // ── 2. tracker features for this frame ─────────────────────
362        let feats = {
363            let _p = crate::cpuprof::spans::span("vid.feats");
364            net.feats(&out.trunk_feats)?
365        };
366
367        // ── 3. propagate the existing objects ──────────────────────
368        let pre_ids = s.obj_ids.clone();
369        let mut trk_masks: Vec<Vec<f32>> = Vec::new();
370        let mut trk_scores: Vec<f32> = Vec::new();
371        if !pre_ids.is_empty() {
372            let _pp = crate::cpuprof::spans::span("vid.propagate");
373            for &id in &pre_ids {
374                let (o, _high) = self.propagate(&net, s, id, frame_idx, &feats)?;
375                trk_scores.push(o.obj_score);
376                trk_masks.push(o.pred_masks.clone());
377                self.store(s, id, frame_idx, false, o);
378            }
379            drop(_pp);
380            // Propagation deliberately does NOT encode memory here. The HF
381            // pipeline defers it until planning has applied its constraints;
382            // `plan` below owns that encode (and execution may replace it when
383            // a new detection causes the tracker forward to run again).
384            fill_holes(&mut trk_masks, self.cfg.fill_hole_area);
385            for (i, &id) in pre_ids.iter().enumerate() {
386                let v = sigm(trk_scores[i]);
387                s.trk_score_fw.entry(id).or_default().insert(frame_idx, v);
388            }
389        }
390
391        // ── 4. planning phase ─────────────────────────────────────
392        let plan = {
393            let _p = crate::cpuprof::spans::span("vid.plan");
394            self.plan(&net, s, frame_idx, &dets, &trk_masks, &feats)?
395        };
396
397        // ── 5. execution phase: new objects become conditioning ───
398        if !plan.new_inds.is_empty() {
399            let _pe = crate::cpuprof::spans::span("vid.execute");
400            let mut pending: Vec<(i64, ObjOut, Vec<f32>)> = Vec::new();
401            for (k, &di) in plan.new_inds.iter().enumerate() {
402                let id = plan.new_ids[k];
403                let mask: Vec<f32> = dets[di]
404                    .mask
405                    .iter()
406                    .map(|&v| if v > 0.0 { 1.0 } else { 0.0 })
407                    .collect();
408                let (o, high) = self.use_mask_as_output(&net, &feats, &mask)?;
409                self.store(s, id, frame_idx, true, o.clone());
410                if !s.obj_ids.contains(&id) {
411                    s.obj_ids.push(id);
412                }
413                pending.push((id, o, high));
414            }
415            // the reference tracker forward re-runs every object that has no
416            // conditioning output at this frame (reconditioned ones are reused)
417            for &id in &pre_ids {
418                if plan.recond_ids.contains(&id) {
419                    continue;
420                }
421                let (o, high) = self.propagate(&net, s, id, frame_idx, &feats)?;
422                self.store(s, id, frame_idx, false, o.clone());
423                pending.push((id, o, high));
424            }
425            let any_mask_input = !pending.is_empty();
426            self.encode_into(&net, s, frame_idx, pending, &feats.l2, IMG, IMG, any_mask_input)?;
427            drop(_pe);
428        }
429
430        // ── 6. outputs ────────────────────────────────────────────
431        let tracks = {
432            let _p = crate::cpuprof::spans::span("vid.outputs");
433            self.build_outputs(s, frame_idx, &pre_ids, &dets, &plan, frame_size)?
434        };
435        let trk_ms = t_trk.elapsed().as_secs_f64() * 1e3;
436
437        for id in &plan.removed {
438            s.obj_ids.retain(|x| x != id);
439            s.objs.remove(id);
440        }
441        s.suppressed.extend(plan.suppressed.iter().copied());
442        s.removed.extend(plan.removed.iter().copied());
443
444        Ok(FrameOut {
445            frame: frame_idx,
446            tracks,
447            detections: dets.len(),
448            presence,
449            det_ms,
450            trk_ms,
451            total_ms: t0.elapsed().as_secs_f64() * 1e3,
452        })
453    }
454
455    // ── detector post-processing ─────────────────────────────────
456
457    /// Scores → threshold → mask-IoU NMS → detections.
458    fn detect(&self, out: &Sam3Out) -> Result<Vec<Det>> {
459        anyhow::ensure!(
460            !out.pred_boxes.is_empty() && out.pred_logits.len() == out.pred_boxes.len() / 4,
461            "forward returned no scores/masks (SAM3_ONLY debug env set?)"
462        );
463        let presence = sigm(out.presence_logit);
464        let nq = out.pred_boxes.len() / 4;
465        let hw = out.pred_masks.len() / nq.max(1);
466        anyhow::ensure!(hw == MASK_SIDE * MASK_SIDE, "unexpected mask shape {hw}");
467        let thr = self.cfg.score_threshold_detection;
468        let mut scores: Vec<f32> = out.pred_logits.iter().map(|&l| sigm(l) * presence).collect();
469        let keep_inital: Vec<bool> = scores.iter().map(|&s| s > thr).collect();
470        let mut keep = keep_inital.clone();
471
472        if self.cfg.det_nms_thresh > 0.0 {
473            let valid: Vec<usize> = (0..nq).filter(|&i| keep[i]).collect();
474            if valid.len() > 1 {
475                let bin = |q: usize| -> Vec<u8> {
476                    out.pred_masks[q * hw..(q + 1) * hw]
477                        .iter()
478                        .map(|&v| u8::from(v > 0.0))
479                        .collect()
480                };
481                let bins: Vec<(usize, Vec<u8>)> = valid.iter().map(|&q| (q, bin(q))).collect();
482                let mut order: Vec<usize> = (0..bins.len()).collect();
483                order.sort_by(|a, b| {
484                    scores[bins[*a].0]
485                        .partial_cmp(&scores[bins[*b].0])
486                        .unwrap_or(std::cmp::Ordering::Equal)
487                });
488                let mut alive = vec![true; bins.len()];
489                for (k, &i) in order.iter().enumerate() {
490                    if !alive[i] {
491                        continue;
492                    }
493                    for &j in &order[k + 1..] {
494                        if alive[j] && mask_iou(&bins[i].1, &bins[j].1) >= self.cfg.det_nms_thresh {
495                            alive[j] = false;
496                        }
497                    }
498                }
499                for (idx, &(q, _)) in bins.iter().enumerate() {
500                    if !alive[idx] {
501                        keep[q] = false;
502                        scores[q] = 0.0;
503                    }
504                }
505            }
506        }
507        let mut dets = Vec::new();
508        for q in 0..nq {
509            if !keep[q] {
510                continue;
511            }
512            dets.push(Det {
513                score: scores[q],
514                mask: out.pred_masks[q * hw..(q + 1) * hw].to_vec(),
515            });
516        }
517        Ok(dets)
518    }
519
520    // ── tracker: one object, one frame ───────────────────────────
521
522    /// Memory-conditioned single-frame inference for an existing object;
523    /// returns the stored output plus the 1008² mask logits for the memory
524    /// encoder.
525    fn propagate(
526        &self,
527        net: &Net<'_>,
528        s: &Session,
529        id: i64,
530        frame_idx: usize,
531        feats: &TrkFeats,
532    ) -> Result<(ObjOut, Vec<f32>)> {
533        let max_ptr = if s.num_frames > 0 {
534            s.num_frames.min(self.cfg.max_object_pointers)
535        } else {
536            self.cfg.max_object_pointers
537        };
538        let st = s.objs.get(&id).ok_or_else(|| anyhow!("unknown object {id}"))?;
539        let cond = self.select_cond(st, frame_idx);
540        let memories = self.gather_memories(st, frame_idx, &cond)?;
541        let pointers = self.gather_pointers(st, frame_idx, &cond);
542        let pix = Net::pix_rows(feats);
543        // memory attention reads the neck's sine PE, the mask decoder the
544        // prompt encoder's Gaussian PE
545        let neck_pos = Net::neck_pos();
546        let pix = {
547            let _p = crate::cpuprof::spans::span("vid.mem_cond");
548            self.memory_cond(net, &pix, &neck_pos, &memories, &pointers, max_ptr)?
549        };
550        let pos = net.image_pos()?;
551
552        let dense_p = net.no_mask_prompt()?;
553        let sparse = net.empty_sparse()?;
554        let dbg = std::env::var("SAM3_VID_DEBUG").is_ok();
555        if dbg {
556            let norm = |v: &[f32]| v.iter().map(|x| x * x).sum::<f32>().sqrt();
557            let mn = pix.iter().cloned().fold(f32::INFINITY, f32::min);
558            let mx = pix.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
559            eprintln!(
560                "[vid] f{frame_idx} obj{id}: memories={} ptr={} pix n={:.3} ({}..{:.3})",
561                memories.len(),
562                pointers.len(),
563                norm(&pix),
564                mn,
565                mx
566            );
567        }
568        let (masks, iou, obj_score, tokens) = {
569            let _p = crate::cpuprof::spans::span("vid.dec");
570            net.mask_decoder(
571                &sam3::rows_to_chw(&pix, EMB, EMB),
572                &pos,
573                &dense_p,
574                &sparse,
575                (&feats.l0, &feats.l1),
576                self.cfg.multimask,
577            )?
578        };
579        let n_masks = masks.len() / (MASK_SIDE * MASK_SIDE);
580        let best = argmax(&iou).min(n_masks.saturating_sub(1));
581        if dbg {
582            let (mut cmn, mut cmx, mut cpos) = (f32::INFINITY, f32::NEG_INFINITY, 0usize);
583            for &v in masks.iter() {
584                cmn = cmn.min(v);
585                cmx = cmx.max(v);
586                cpos += usize::from(v > 0.0);
587            }
588            let base = best * MASK_SIDE * MASK_SIDE;
589            let spos = masks[base..base + MASK_SIDE * MASK_SIDE].iter().filter(|&&v| v > 0.0).count();
590            eprintln!(
591                "[vid] f{frame_idx} obj{id}: decoder iou={iou:?} obj_score={obj_score:.4} masks {cmn:.3}..{cmx:.3} pos={cpos}/{} chosen pos={spos}",
592                masks.len()
593            );
594        }
595        let low = masks[best * MASK_SIDE * MASK_SIDE..(best + 1) * MASK_SIDE * MASK_SIDE].to_vec();
596        let tok = &tokens[best * HID..(best + 1) * HID];
597        let appearing = obj_score > 0.0;
598        let low = if appearing {
599            low
600        } else {
601            if dbg {
602                eprintln!("[vid] f{frame_idx} obj{id}: NOT appearing (obj_score {obj_score}) -> mask forced to NO_OBJ");
603            }
604            vec![NO_OBJ_SCORE; MASK_SIDE * MASK_SIDE]
605        };
606        if dbg {
607            let pos = low.iter().filter(|&&v| v > 0.0).count();
608            eprintln!("[vid] f{frame_idx} obj{id}: propagate mask pos={pos}/{}", low.len());
609        }
610        let high = {
611            let _p = crate::cpuprof::spans::span("vid.resize");
612            resize2d(&low, MASK_SIDE, MASK_SIDE, IMG, IMG)
613        };
614        let pointer = self.object_pointer(net, tok, appearing)?;
615        Ok((
616            ObjOut { pred_masks: low, obj_score, pointer, mem_feat: None, mem_pos: None },
617            high,
618        ))
619    }
620
621    /// `_use_mask_as_output`: the detection mask IS the output; the decoder
622    /// runs once to produce the object pointer.
623    fn use_mask_as_output(
624        &self,
625        net: &Net<'_>,
626        feats: &TrkFeats,
627        mask: &[f32],
628    ) -> Result<(ObjOut, Vec<f32>)> {
629        let appearing = mask.iter().any(|&v| v > 0.0);
630        let obj_score = if appearing { 10.0 } else { -10.0 };
631        // the 288² binary mask is first resampled to the tracker's 1008²
632        // input, then scaled to logits (*20-10) — `_use_mask_as_output`
633        let hi = resize2d(mask, MASK_SIDE, MASK_SIDE, IMG, IMG);
634        let high: Vec<f32> = hi
635            .iter()
636            .map(|v| v * self.cfg.sigmoid_scale + self.cfg.sigmoid_bias)
637            .collect();
638        let low = resize2d(&high, IMG, IMG, MASK_SIDE, MASK_SIDE);
639
640        // pointer path: 4× downsample of the 1008² mask → 288² prompt → decoder
641        let down = mask_downsample(net, &hi)?;
642        let down = resize2d(&down, IMG / 4, IMG / 4, MASK_SIDE, MASK_SIDE);
643        let dense_p = net.mask_prompt(&down)?;
644        let sparse = net.empty_sparse()?;
645        let pix = Net::pix_rows(feats);
646        let pos = net.image_pos()?;
647        let (_m, iou, dec_score, tokens) = net.mask_decoder(
648            &sam3::rows_to_chw(&pix, EMB, EMB),
649            &pos,
650            &dense_p,
651            &sparse,
652            (&feats.l0, &feats.l1),
653            self.cfg.multimask,
654        )?;
655        let n = tokens.len() / HID;
656        // `_single_frame_forward` uses the same best-IoU token for the
657        // object pointer as for the selected mask (not always the last one).
658        let best = argmax(&iou).min(n.saturating_sub(1));
659        let tok = &tokens[best * HID..(best + 1) * HID];
660        let mut pointer = self.object_pointer(net, tok, dec_score > 0.0)?;
661        if !appearing {
662            pointer = net.v("trk.no_obj_ptr")?;
663        }
664        Ok((
665            ObjOut { pred_masks: low, obj_score, pointer, mem_feat: None, mem_pos: None },
666            high,
667        ))
668    }
669
670    fn object_pointer(&self, net: &Net<'_>, tok: &[f32], appearing: bool) -> Result<Vec<f32>> {
671        let p = "trk.obj_ptr_proj.";
672        // FeedForward(256,256,256,3): proj_in, layers.0, proj_out
673        let mut h = dense(
674            &net.w(&format!("{p}layers.0.weight"))?,
675            &net.v(&format!("{p}layers.0.bias"))?,
676            tok,
677            1,
678            net.pool,
679        );
680        h = relu_vec(&h);
681        h = dense(
682            &net.w(&format!("{p}layers.1.weight"))?,
683            &net.v(&format!("{p}layers.1.bias"))?,
684            &h,
685            1,
686            net.pool,
687        );
688        h = relu_vec(&h);
689        let ptr = dense(
690            &net.w(&format!("{p}layers.2.weight"))?,
691            &net.v(&format!("{p}layers.2.bias"))?,
692            &h,
693            1,
694            net.pool,
695        );
696        if appearing {
697            Ok(ptr)
698        } else {
699            net.v("trk.no_obj_ptr")
700        }
701    }
702
703    // ── memory bank ──────────────────────────────────────────────
704
705    /// Conditioning frames closest to `frame_idx`
706    /// (`_select_closest_cond_frames`).
707    fn select_cond(&self, st: &ObjState, frame_idx: usize) -> Vec<usize> {
708        let max = self.cfg.max_cond_frame_num;
709        if max == 0 || st.cond.len() <= max {
710            return st.cond.keys().copied().collect();
711        }
712        let mut sel: Vec<usize> = Vec::new();
713        if let Some(before) = st.cond.keys().copied().filter(|&t| t < frame_idx).next_back() {
714            sel.push(before);
715        }
716        if let Some(after) = st.cond.keys().copied().filter(|&t| t >= frame_idx).next() {
717            sel.push(after);
718        }
719        let remain = max.saturating_sub(sel.len());
720        let mut rest: Vec<usize> = st.cond.keys().copied().filter(|t| !sel.contains(t)).collect();
721        rest.sort_by_key(|t| (*t as i64 - frame_idx as i64).abs());
722        sel.extend(rest.into_iter().take(remain));
723        sel
724    }
725
726    /// `(offset, mem_feat, mem_pos)` for every memory frame
727    /// (`_gather_memory_frame_outputs` + `_build_memory_attention_inputs`).
728    fn gather_memories(
729        &self,
730        st: &ObjState,
731        frame_idx: usize,
732        cond: &[usize],
733    ) -> Result<Vec<(usize, Vec<f32>, Vec<f32>)>> {
734        let mut out = Vec::new();
735        for &t in cond {
736            if let Some(o) = st.cond.get(&t) {
737                if let (Some(f), Some(p)) = (&o.mem_feat, &o.mem_pos) {
738                    out.push((0, f.clone(), p.clone()));
739                }
740            }
741        }
742        for off in (1..self.cfg.num_maskmem).rev() {
743            let prev = frame_idx as i64 - off as i64;
744            if prev < 0 {
745                continue;
746            }
747            if let Some(o) = st.non_cond.get(&(prev as usize)) {
748                if let (Some(f), Some(p)) = (&o.mem_feat, &o.mem_pos) {
749                    out.push((off, f.clone(), p.clone()));
750                }
751            }
752        }
753        Ok(out)
754    }
755
756    /// Object pointers of past frames (`_get_object_pointers`).
757    fn gather_pointers(&self, st: &ObjState, frame_idx: usize, cond: &[usize]) -> Vec<(usize, Vec<f32>)> {
758        let mut offs = Vec::new();
759        let mut toks = Vec::new();
760        for &t in cond {
761            if t > frame_idx {
762                continue;
763            }
764            if let Some(o) = st.cond.get(&t) {
765                offs.push(frame_idx - t);
766                toks.push(o.pointer.clone());
767            }
768        }
769        for d in 1..self.cfg.max_object_pointers.max(2) {
770            let prev = frame_idx as i64 - d as i64;
771            if prev < 0 {
772                break;
773            }
774            if let Some(o) = st.non_cond.get(&(prev as usize)) {
775                offs.push(d);
776                toks.push(o.pointer.clone());
777            }
778        }
779        offs.into_iter().zip(toks).collect()
780    }
781
782    /// `_prepare_memory_conditioned_features`: current features fused with
783    /// the memory bank through the 4-layer RoPE memory attention.
784    fn memory_cond(
785        &self,
786        net: &Net<'_>,
787        pix: &[f32],
788        pos: &[f32],
789        memories: &[(usize, Vec<f32>, Vec<f32>)],
790        pointers: &[(usize, Vec<f32>)],
791        max_ptr: usize,
792    ) -> Result<Vec<f32>> {
793        let n = EMB * EMB;
794        if memories.is_empty() {
795            // conditioning frame: features + no_memory_embedding (a zero buffer)
796            return Ok(pix.to_vec());
797        }
798        let tpos = net.v("trk.maskmem_tpos_enc")?; // (7, MEM)
799        let n_tpos = tpos.len() / MEM;
800        let mut mem: Vec<f32> = Vec::new();
801        let mut mem_pos: Vec<f32> = Vec::new();
802        for (off, f, p) in memories {
803            // cond frames (offset 0) index the LAST temporal row — the
804            // reference's `tpos[offset - 1]` with negative indexing.
805            let row = if *off == 0 { n_tpos - 1 } else { (*off - 1).min(n_tpos - 1) };
806            for (i, v) in p.iter().enumerate() {
807                mem_pos.push(v + tpos[row * MEM + i % MEM]);
808            }
809            mem.extend_from_slice(f);
810        }
811        if !pointers.is_empty() {
812            let max_diff = max_ptr.saturating_sub(1).max(1) as f32;
813            let offs: Vec<f32> = pointers.iter().map(|(o, _)| *o as f32 / max_diff).collect();
814            let pe = sine_1d(&offs, HID);
815            let tw = net.w("trk.obj_ptr_tpos_proj.weight")?;
816            let tb = net.v("trk.obj_ptr_tpos_proj.bias")?;
817            for (i, (_o, p)) in pointers.iter().enumerate() {
818                let proj = dense(&tw, &tb, &pe[i * HID..(i + 1) * HID], 1, net.pool);
819                let splits = HID / MEM;
820                for s in 0..splits {
821                    mem.extend_from_slice(&p[s * MEM..(s + 1) * MEM]);
822                    mem_pos.extend_from_slice(&proj);
823                }
824            }
825        }
826        // Object pointers are keys/values too: their four 64-wide slices are
827        // appended after the spatial memory and excluded only from RoPE.
828        let mem_len = mem.len() / MEM;
829
830        // RoPE over the 72×72 grid: dim 256, spatial half 128 → 64 freqs
831        let head_dim = HID;
832        let spatial = head_dim / 2;
833        let mut inv = vec![0f32; spatial / 2];
834        for i in 0..spatial / 2 {
835            inv[i] = 1.0 / 10000f32.powf((2 * i) as f32 / spatial as f32);
836        }
837        let mut cos = vec![0f32; n * head_dim];
838        let mut sin = vec![0f32; n * head_dim];
839        for y in 0..EMB {
840            for x in 0..EMB {
841                let base = (y * EMB + x) * head_dim;
842                // position_ids row 0 = w (x), row 1 = h (y)
843                for (k, src) in [x, y].iter().enumerate() {
844                    for (j, iv) in inv.iter().enumerate() {
845                        let ang = *src as f32 * *iv;
846                        let dst = base + k * spatial + j * 2; // repeat_interleave(2)
847                        cos[dst] = ang.cos();
848                        cos[dst + 1] = ang.cos();
849                        sin[dst] = ang.sin();
850                        sin[dst + 1] = ang.sin();
851                    }
852                }
853            }
854        }
855
856        let mut q = pix.to_vec();
857        for (a, b) in q.iter_mut().zip(pos) {
858            *a += 0.1 * *b;
859        }
860        let dump_dir = std::env::var("SAM3_VID_DUMP").ok();
861        let dump = |name: &str, v: &[f32]| {
862            if let Some(d) = &dump_dir {
863                let _ = std::fs::create_dir_all(d);
864                let _ = std::fs::write(format!("{d}/{name}.f32"), Self::f32_bytes(v));
865            }
866        };
867        dump("initial_seq", &q);
868        dump("cos", &cos);
869        dump("sin", &sin);
870        dump("memory", &mem);
871        dump("memory_pos", &mem_pos);
872        let tb = "trk.transformer.encoder.";
873        for layer in 0..4 {
874            let lb = format!("{tb}layers.{layer}.");
875            let mut h = q.clone();
876            net.ln_rows(&mut h, &format!("{lb}norm1.weight"), &format!("{lb}norm1.bias"), 1e-5)?;
877            dump(&format!("l{layer}_norm1"), &h);
878            let mut ao = vec![0f32; n * HID];
879            net.attn(&format!("{lb}self_attn."), &h, &h, &h, n, n, HID, HID, HID, Some((&cos, &sin, n)), &mut ao)?;
880            dump(&format!("l{layer}_self"), &ao);
881            for (a, c) in q.iter_mut().zip(&ao) {
882                *a += *c;
883            }
884            dump(&format!("l{layer}_res1"), &q);
885            let mut h = q.clone();
886            net.ln_rows(&mut h, &format!("{lb}norm2.weight"), &format!("{lb}norm2.bias"), 1e-5)?;
887            dump(&format!("l{layer}_norm2"), &h);
888            // keys = memory + memory position encoding (kv_in = MEM)
889            let k: Vec<f32> = mem
890                .chunks(MEM)
891                .zip(mem_pos.chunks(MEM))
892                .flat_map(|(a, b)| a.iter().zip(b).map(|(x, y)| x + y))
893                .collect();
894            if layer == 0 {
895                dump("cross_key", &k);
896            }
897            let mut ao = vec![0f32; n * HID];
898            net.attn(
899                &format!("{lb}cross_attn_image."),
900                &h,
901                &k,
902                &mem,
903                n,
904                mem_len,
905                HID,
906                MEM,
907                HID,
908                Some((&cos, &sin, mem_len - pointers.len() * (HID / MEM))),
909                &mut ao,
910            )?;
911            dump(&format!("l{layer}_cross"), &ao);
912            for (a, c) in q.iter_mut().zip(&ao) {
913                *a += *c;
914            }
915            dump(&format!("l{layer}_res2"), &q);
916            let mut h = q.clone();
917            net.ln_rows(&mut h, &format!("{lb}norm3.weight"), &format!("{lb}norm3.bias"), 1e-5)?;
918            let a = net.lin(&format!("{lb}"), "linear1", &h, n)?;
919            let a = relu_vec(&a);
920            let a = net.lin(&format!("{lb}"), "linear2", &a, n)?;
921            for (x, c) in q.iter_mut().zip(&a) {
922                *x += *c;
923            }
924            dump(&format!("l{layer}_out"), &q);
925        }
926        net.ln_rows(&mut q, &format!("{tb}norm.weight"), &format!("{tb}norm.bias"), 1e-5)?;
927        dump("final", &q);
928        Ok(q)
929    }
930
931    /// Memory encoder: mask → stride-16 downsampler, fused with the pixel
932    /// features, ConvNeXt fuser, projection to `mem_dim` + sine PE.
933    #[allow(dead_code)]
934    fn f32_bytes(v: &[f32]) -> Vec<u8> {
935        let mut b = Vec::with_capacity(v.len() * 4);
936        for &x in v {
937            b.extend_from_slice(&x.to_le_bytes());
938        }
939        b
940    }
941
942    fn encode_memory(
943        &self,
944        net: &Net<'_>,
945        pix: &[f32],        // (256,72,72) CHW
946        mask_logits: &[f32], // HxW logits
947        h: usize,
948        w: usize,
949        obj_score: f32,
950        from_pts: bool,
951    ) -> Result<(Vec<f32>, Vec<f32>)> {
952        let mut m = resize2d(mask_logits, h, w, MEM_SIDE, MEM_SIDE);
953        if from_pts {
954            for v in m.iter_mut() {
955                *v = if *v > 0.0 { 1.0 } else { 0.0 };
956            }
957        } else {
958            for v in m.iter_mut() {
959                *v = sigm(*v);
960            }
961        }
962        for v in m.iter_mut() {
963            *v = *v * self.cfg.sigmoid_scale + self.cfg.sigmoid_bias;
964        }
965        let dbg = std::env::var("SAM3_VID_DEBUG").is_ok();
966        if dbg {
967            let mn = m.iter().cloned().fold(f32::INFINITY, f32::min);
968            let mx = m.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
969            let pos = m.iter().filter(|&&v| v > 0.0).count();
970            eprintln!("[mem] mask input ({}): {mn:.3}..{mx:.3} pos={}/{} scale={} bias={}", m.len(), pos, m.len(), self.cfg.sigmoid_scale, self.cfg.sigmoid_bias);
971        }
972        // mask downsampler: 4 × (conv3 s2 + LN + GELU), then 1×1
973        let p = "trk.maskmem_backbone.mask_downsampler.encoder.";
974        let mut x = m;
975        let mut side = MEM_SIDE;
976        let mut c_in = 1usize;
977        for step in [0usize, 3, 6, 9] {
978            let wgt = net.v(&format!("{p}{step}.weight"))?;
979            let b = net.v(&format!("{p}{step}.bias"))?;
980            x = sam3::conv2d(&wgt, &b, &x, c_in, side, side, 3, 3, 2, 1, net.pool);
981            side /= 2;
982            let ch = b.len();
983            let lw = net.v(&format!("{p}{}.weight", step + 1))?;
984            let lb = net.v(&format!("{p}{}.bias", step + 1))?;
985            ln_chw(&mut x, ch, side * side, &lw, &lb, 1e-6);
986            x.iter_mut().for_each(|v| *v = gelu_erf(*v));
987            c_in = ch;
988            if dbg {
989                let mn = x.iter().cloned().fold(f32::INFINITY, f32::min);
990                let mx = x.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
991                eprintln!("[mem] after layer {step}: side={side} ch={ch} {mn:.3}..{mx:.3}");
992            }
993        }
994        let w12 = net.v(&format!("{p}12.weight"))?;
995        let b12 = net.v(&format!("{p}12.bias"))?;
996        x = sam3::conv2d(&w12, &b12, &x, c_in, side, side, 1, 1, 1, 0, net.pool);
997        anyhow::ensure!(x.len() == HID * EMB * EMB, "mask downsampler output {}", x.len());
998        if dbg {
999            let rms = |v: &[f32]| v.iter().map(|x| x * x).sum::<f32>().sqrt() / (v.len() as f32).sqrt();
1000            eprintln!("[mem] mask branch rms={:.4} | pix rms={:.4}", rms(&x), rms(pix));
1001        }
1002
1003        // pixel projection + mask
1004        let pw = net.v("trk.maskmem_backbone.pix_feat_proj.weight")?;
1005        let pb = net.v("trk.maskmem_backbone.pix_feat_proj.bias")?;
1006        let mut f = sam3::conv2d(&pw, &pb, pix, HID, EMB, EMB, 1, 1, 1, 0, net.pool);
1007        if dbg {
1008            let rms = |v: &[f32]| v.iter().map(|x| x * x).sum::<f32>().sqrt() / (v.len() as f32).sqrt();
1009            eprintln!("[mem] pix_proj rms={:.4}", rms(&f));
1010        }
1011        for (a, b) in f.iter_mut().zip(&x) {
1012            *a += *b;
1013        }
1014        if dbg {
1015            let rms = |v: &[f32]| v.iter().map(|x| x * x).sum::<f32>().sqrt() / (v.len() as f32).sqrt();
1016            eprintln!("[mem] fused pre-fuser rms={:.4}", rms(&f));
1017        }
1018        let dump_dir = std::env::var("SAM3_VID_DUMP").ok();
1019        if let Some(d) = &dump_dir {
1020            let _ = std::fs::create_dir_all(d);
1021            let _ = std::fs::write(format!("{d}/mask_branch.f32"), Self::f32_bytes(&x));
1022            let _ = std::fs::write(format!("{d}/pix_proj.f32"), Self::f32_bytes(&f));
1023        }
1024        // ConvNeXt fuser ×2 — `out = block_in + gamma * pwconv2(gelu(pwconv1(LN(dwconv(block_in)))))`
1025        for layer in 0..2 {
1026            let lp = format!("trk.maskmem_backbone.fuser.layers.{layer}.");
1027            let block_in = f.clone(); // the residual branch springs from the BLOCK INPUT
1028            let dw = net.v(&format!("{lp}dwconv.weight"))?;
1029            let db = net.v(&format!("{lp}dwconv.bias"))?;
1030            f = depthwise(&dw, &db, &f, HID, EMB, EMB, 7, net.pool);
1031            let nw = net.v(&format!("{lp}norm.weight"))?;
1032            let nb = net.v(&format!("{lp}norm.bias"))?;
1033            ln_chw(&mut f, HID, EMB * EMB, &nw, &nb, 1e-6);
1034            let ln_out = if dump_dir.is_some() { Some(f.clone()) } else { None };
1035            // ConvNeXt pointwise convs run over the NHWC layout: rows =
1036            // spatial positions, columns = channels
1037            let rows = EMB * EMB;
1038            let dw_out = if dump_dir.is_some() { Some(f.clone()) } else { None };
1039            let w1 = net.w(&format!("{lp}pwconv1.weight"))?;
1040            let b1 = net.v(&format!("{lp}pwconv1.bias"))?;
1041            let h_rows = sam3::chw_to_rows(&f, EMB, EMB);
1042            let mut hh = dense(&w1, &b1, &h_rows, rows, net.pool); // (HW, 1024)
1043            hh.iter_mut().for_each(|v| *v = gelu_erf(*v));
1044            let pw1_out = if dump_dir.is_some() { Some(hh.clone()) } else { None };
1045            let w2 = net.w(&format!("{lp}pwconv2.weight"))?;
1046            let b2 = net.v(&format!("{lp}pwconv2.bias"))?;
1047            let hh = dense(&w2, &b2, &hh, rows, net.pool); // (HW, 256)
1048            let g = net.v(&format!("{lp}gamma"))?;
1049            if let Some(d) = &dump_dir {
1050                let _ = std::fs::write(format!("{d}/fuser_gamma{layer}.f32"), Self::f32_bytes(&g));
1051            }
1052            for c in 0..HID {
1053                for i in 0..rows {
1054                    f[c * rows + i] = block_in[c * rows + i] + g[c] * hh[i * HID + c];
1055                }
1056            }
1057            if let Some(d) = &dump_dir {
1058                let _ = std::fs::write(format!("{d}/fuser_layer{layer}.f32"), Self::f32_bytes(&f));
1059                if layer == 0 {
1060                    if let Some(v) = &dw_out {
1061                        let _ = std::fs::write(format!("{d}/f0_dwconv.f32"), Self::f32_bytes(v));
1062                    }
1063                    if let Some(v) = &ln_out {
1064                        let _ = std::fs::write(format!("{d}/f0_norm.f32"), Self::f32_bytes(v));
1065                    }
1066                    if let Some(v) = &pw1_out {
1067                        let _ = std::fs::write(format!("{d}/f0_pw1.f32"), Self::f32_bytes(v));
1068                    }
1069                    let _ = std::fs::write(format!("{d}/f0_pw2.f32"), Self::f32_bytes(&hh));
1070                }
1071            }
1072        }
1073        if std::env::var("SAM3_VID_DEBUG").is_ok() {
1074            let rms = |v: &[f32]| v.iter().map(|x| x * x).sum::<f32>().sqrt() / (v.len() as f32).sqrt();
1075            eprintln!("[mem] post-fuser rms={:.4}", rms(&f));
1076        }
1077        let ow = net.v("trk.maskmem_backbone.out_proj.weight")?;
1078        let ob = net.v("trk.maskmem_backbone.out_proj.bias")?;
1079        let mut f = sam3::conv2d(&ow, &ob, &f, HID, EMB, EMB, 1, 1, 1, 0, net.pool);
1080        anyhow::ensure!(f.len() == MEM * EMB * EMB, "memory projection {}", f.len());
1081        if self.cfg.occlusion_spatial && obj_score <= 0.0 {
1082            let e = net.v("trk.no_obj_embed_spatial")?;
1083            let hw = EMB * EMB;
1084            for i in 0..hw {
1085                for (c, ev) in e.iter().enumerate() {
1086                    f[c * hw + i] += ev;
1087                }
1088            }
1089        }
1090        let pos = sine(EMB, EMB, MEM / 2);
1091        let rows = sam3::chw_to_rows(&f, EMB, EMB);
1092        Ok((rows, pos))
1093    }
1094
1095    // ── storage helpers ──────────────────────────────────────────
1096
1097    fn store(&self, s: &mut Session, id: i64, frame_idx: usize, cond: bool, o: ObjOut) {
1098        let st = s.objs.entry(id).or_default();
1099        if cond {
1100            st.cond.insert(frame_idx, o);
1101        } else {
1102            st.non_cond.insert(frame_idx, o);
1103        }
1104    }
1105
1106    /// Encode memories for the outputs just written
1107    /// (`_batch_encode_memories`).
1108    fn encode_into(
1109        &self,
1110        net: &Net<'_>,
1111        s: &mut Session,
1112        frame_idx: usize,
1113        pending: Vec<(i64, ObjOut, Vec<f32>)>,
1114        pix: &[f32],
1115        h: usize,
1116        w: usize,
1117        from_pts: bool,
1118    ) -> Result<()> {
1119        if pending.is_empty() {
1120            return Ok(());
1121        }
1122        let _p = crate::cpuprof::spans::span("vid.encode");
1123        for (id, o, high) in pending {
1124            let (mf, mp) = self.encode_memory(net, pix, &high, h, w, o.obj_score, from_pts)?;
1125            let st = s.objs.entry(id).or_default();
1126            let slot = if st.cond.contains_key(&frame_idx) {
1127                st.cond.get_mut(&frame_idx)
1128            } else {
1129                st.non_cond.get_mut(&frame_idx)
1130            };
1131            if let Some(slot) = slot {
1132                slot.mem_feat = Some(mf);
1133                slot.mem_pos = Some(mp);
1134            }
1135        }
1136        Ok(())
1137    }
1138}
1139
1140// ───────────────────── planning ─────────────────────
1141
1142struct Plan {
1143    new_inds: Vec<usize>,
1144    new_ids: Vec<i64>,
1145    recond_ids: HashSet<i64>,
1146    removed: HashSet<i64>,
1147    suppressed: HashSet<i64>,
1148}
1149
1150impl Sam3Video {
1151    #[allow(clippy::too_many_arguments)]
1152    fn plan(
1153        &self,
1154        net: &Net<'_>,
1155        s: &mut Session,
1156        frame_idx: usize,
1157        dets: &[Det],
1158        trk_masks: &[Vec<f32>],
1159        feats: &TrkFeats,
1160    ) -> Result<Plan> {
1161        let cfg = &self.cfg;
1162        let pre_ids: Vec<i64> = s.obj_ids.clone();
1163        let det_masks: Vec<&[f32]> = dets.iter().map(|d| d.mask.as_slice()).collect();
1164        let det_scores: Vec<f32> = dets.iter().map(|d| d.score).collect();
1165        let assoc = associate(&det_masks, &det_scores, trk_masks, &pre_ids, cfg);
1166        if std::env::var("SAM3_VID_DEBUG").is_ok() {
1167            let trk_pos: Vec<usize> = trk_masks.iter().map(|m| m.iter().filter(|&&v| v > 0.0).count()).collect();
1168            let det_pos: Vec<usize> = det_masks.iter().map(|m| m.iter().filter(|&&v| v > 0.0).count()).collect();
1169            eprintln!(
1170                "[vid] f{frame_idx} plan: dets={} (pos {:?}, scores {:?}) trk={} (pos {:?}) new={:?} unmatched={:?} empty={:?}",
1171                det_masks.len(), det_pos, det_scores, trk_masks.len(), trk_pos,
1172                assoc.new_inds, assoc.unmatched, assoc.empty
1173            );
1174        }
1175
1176        // object-id budget
1177        let mut new_inds = assoc.new_inds;
1178        if pre_ids.len() + new_inds.len() > cfg.max_num_objects {
1179            let room = cfg.max_num_objects.saturating_sub(pre_ids.len());
1180            let mut by_score = new_inds.clone();
1181            by_score.sort_by(|a, b| {
1182                det_scores[*b].partial_cmp(&det_scores[*a]).unwrap_or(std::cmp::Ordering::Equal)
1183            });
1184            new_inds = by_score.into_iter().take(room).collect();
1185            new_inds.sort_unstable();
1186        }
1187        let start = s.max_obj_id + 1;
1188        let new_ids: Vec<i64> = (0..new_inds.len() as i64).map(|k| start + k).collect();
1189
1190        // ── hotstart (`_process_hotstart`) ─────────────────────────
1191        let mut removed: HashSet<i64> = HashSet::new();
1192        let mut suppressed: HashSet<i64> = HashSet::new();
1193        for (k, &di) in new_inds.iter().enumerate() {
1194            let id = new_ids[k];
1195            let _ = di;
1196            s.first_frame.entry(id).or_insert(frame_idx as i64);
1197            s.keep_alive.insert(id, cfg.init_trk_keep_alive);
1198        }
1199        let matched: HashSet<i64> = assoc.det_to_matched.values().flatten().copied().collect();
1200        for id in &matched {
1201            let ka = s.keep_alive.get(id).copied().unwrap_or(cfg.init_trk_keep_alive);
1202            s.keep_alive.insert(*id, cfg.max_trk_keep_alive.max(ka + 1));
1203        }
1204        for id in &assoc.unmatched {
1205            s.unmatched.entry(*id).or_default().push(frame_idx);
1206            let ka = s.keep_alive.get(id).copied().unwrap_or(cfg.init_trk_keep_alive);
1207            s.keep_alive.insert(*id, cfg.min_trk_keep_alive.max(ka - 1));
1208        }
1209        if cfg.decrease_keep_alive_for_empty {
1210            for id in &assoc.empty {
1211                let ka = s.keep_alive.get(id).copied().unwrap_or(cfg.init_trk_keep_alive);
1212                s.keep_alive.insert(*id, cfg.min_trk_keep_alive.max(ka - 1));
1213            }
1214        }
1215        let hotstart_diff = frame_idx as i64 - cfg.hotstart_delay;
1216        let already_removed: HashSet<i64> = s.removed.clone();
1217        for (id, frames) in s.unmatched.clone() {
1218            if already_removed.contains(&id) || removed.contains(&id) {
1219                continue;
1220            }
1221            if frames.len() >= cfg.hotstart_unmatch_thresh {
1222                if s.first_frame.get(&id).copied().unwrap_or(0) > hotstart_diff {
1223                    removed.insert(id);
1224                }
1225            }
1226            let ka = s.keep_alive.get(&id).copied().unwrap_or(0);
1227            if ka <= 0
1228                && !cfg.suppress_unmatched_only_within_hotstart
1229                && !already_removed.contains(&id)
1230                && !removed.contains(&id)
1231            {
1232                suppressed.insert(id);
1233            }
1234        }
1235        for ids in assoc.det_to_matched.values() {
1236            if ids.len() < 2 {
1237                continue;
1238            }
1239            let first = *ids
1240                .iter()
1241                .min_by_key(|i| s.first_frame.get(i).copied().unwrap_or(i64::MAX))
1242                .unwrap();
1243            for id in ids {
1244                if *id != first {
1245                    s.overlap.entry((first, *id)).or_default().push(frame_idx);
1246                }
1247            }
1248        }
1249        let mut pairs: Vec<((i64, i64), Vec<usize>)> = s.overlap.clone().into_iter().collect();
1250        pairs.sort();
1251        for ((_, id), frames) in pairs {
1252            if s.removed.contains(&id) || removed.contains(&id) {
1253                continue;
1254            }
1255            if s.first_frame.get(&id).copied().unwrap_or(0) > hotstart_diff
1256                && frames.len() >= cfg.hotstart_dup_thresh
1257            {
1258                removed.insert(id);
1259            }
1260        }
1261        s.removed.extend(removed.iter().copied());
1262
1263        // ── periodic reconditioning (`_prepare_recondition_masks`) ──
1264        let mut recond_ids: HashSet<i64> = HashSet::new();
1265        let mut recond_masks: Vec<(usize, Vec<f32>)> = Vec::new();
1266        if cfg.recondition_every_nth_frame > 0
1267            && frame_idx as i64 % cfg.recondition_every_nth_frame == 0
1268            && !assoc.trk_to_high_conf.is_empty()
1269        {
1270            for (&id, &di) in &assoc.trk_to_high_conf {
1271                let Some(i) = pre_ids.iter().position(|x| *x == id) else { continue };
1272                let trk_score = s
1273                    .trk_score_fw
1274                    .get(&id)
1275                    .and_then(|m| m.get(&frame_idx))
1276                    .copied()
1277                    .unwrap_or(0.0);
1278                if trk_score <= cfg.high_conf_thresh {
1279                    continue;
1280                }
1281                recond_ids.insert(id);
1282                if cfg.recondition_on_trk_masks {
1283                    recond_masks.push((i, trk_masks[i].clone()));
1284                } else {
1285                    recond_masks.push((
1286                        i,
1287                        dets[di].mask.iter().map(|&v| if v > 0.0 { 1.0 } else { 0.0 }).collect(),
1288                    ));
1289                }
1290            }
1291        }
1292
1293        // ── recent-occlusion suppression ────────────────────────────
1294        let mut trk = trk_masks.to_vec();
1295        if cfg.suppress_recent_occlusion_thresh > 0.0 && !trk.is_empty() {
1296            let bins: Vec<Vec<u8>> = trk
1297                .iter()
1298                .map(|m| m.iter().map(|&v| u8::from(v > 0.0)).collect())
1299                .collect();
1300            let n = bins.len();
1301            let mut last: Vec<i64> = pre_ids
1302                .iter()
1303                .map(|id| {
1304                    s.last_occluded.get(id).copied().unwrap_or(if s.removed.contains(id) {
1305                        100_000
1306                    } else {
1307                        -1
1308                    })
1309                })
1310                .collect();
1311            let mut to_suppress = vec![false; n];
1312            if n > 1 {
1313                for i in 0..n {
1314                    for j in (i + 1)..n {
1315                        if mask_iou(&bins[i], &bins[j]) < cfg.suppress_recent_occlusion_thresh {
1316                            continue;
1317                        }
1318                        if last[i] > last[j] && last[j] > -1 {
1319                            to_suppress[i] = true;
1320                        }
1321                        if last[j] > last[i] && last[i] > -1 {
1322                            to_suppress[j] = true;
1323                        }
1324                    }
1325                }
1326            }
1327            for (i, m) in trk.iter_mut().enumerate() {
1328                if bins[i].iter().all(|&v| v == 0) || to_suppress[i] {
1329                    for v in m.iter_mut() {
1330                        *v = -10.0;
1331                    }
1332                    last[i] = frame_idx as i64;
1333                }
1334            }
1335            for (i, id) in pre_ids.iter().enumerate() {
1336                s.last_occluded.insert(*id, last[i]);
1337            }
1338        }
1339
1340        // ── `_tracker_update_memories`: plan → constraints → encode ─
1341        let mut per_obj = trk;
1342        for (i, m) in &recond_masks {
1343            per_obj[*i] = m.clone();
1344        }
1345        for (i, &id) in pre_ids.iter().enumerate() {
1346            if recond_ids.contains(&id) {
1347                if let Some(st) = s.objs.get_mut(&id) {
1348                    if let Some(o) = st.non_cond.remove(&frame_idx) {
1349                        st.cond.insert(frame_idx, o);
1350                    }
1351                }
1352            }
1353        }
1354        let per_obj = suppress_area_shrink(&per_obj);
1355        if !per_obj.is_empty() {
1356            let mut pending = Vec::new();
1357            for (i, &id) in pre_ids.iter().enumerate() {
1358                let area = per_obj[i].iter().any(|&v| v > 0.0);
1359                let score = if area { 10.0 } else { -10.0 };
1360                if let Some(st) = s.objs.get_mut(&id) {
1361                    let slot = if st.cond.contains_key(&frame_idx) {
1362                        st.cond.get_mut(&frame_idx)
1363                    } else {
1364                        st.non_cond.get_mut(&frame_idx)
1365                    };
1366                    if let Some(slot) = slot {
1367                        slot.pred_masks = per_obj[i].clone();
1368                        slot.obj_score = score;
1369                    }
1370                }
1371                pending.push((
1372                    id,
1373                    ObjOut {
1374                        pred_masks: per_obj[i].clone(),
1375                        obj_score: score,
1376                        pointer: Vec::new(),
1377                        mem_feat: None,
1378                        mem_pos: None,
1379                    },
1380                    per_obj[i].clone(),
1381                ));
1382            }
1383            self.encode_into(net, s, frame_idx, pending, &feats.l2, MASK_SIDE, MASK_SIDE, false)?;
1384        }
1385
1386        // ── metadata updates ───────────────────────────────────────
1387        for (k, &id) in new_ids.iter().enumerate() {
1388            s.score.insert(id, det_scores[new_inds[k]]);
1389            s.trk_score_fw.entry(id).or_default().insert(frame_idx, det_scores[new_inds[k]]);
1390        }
1391        for id in &removed {
1392            s.score.insert(*id, -1e4);
1393            s.trk_score_fw.entry(*id).or_default().insert(frame_idx, -1e4);
1394        }
1395        s.max_obj_id = (start + new_ids.len() as i64 - 1).max(s.max_obj_id);
1396
1397        Ok(Plan { new_inds, new_ids, recond_ids, removed, suppressed })
1398    }
1399}
1400
1401// ───────────────────── association ─────────────────────
1402
1403struct Assoc {
1404    new_inds: Vec<usize>,
1405    unmatched: Vec<i64>,
1406    empty: Vec<i64>,
1407    det_to_matched: HashMap<usize, Vec<i64>>,
1408    trk_to_high_conf: HashMap<i64, usize>,
1409}
1410
1411/// `_associate_det_trk`: detections against the propagated masklets.
1412fn associate(
1413    det_masks: &[&[f32]],
1414    det_scores: &[f32],
1415    trk_masks: &[Vec<f32>],
1416    trk_ids: &[i64],
1417    cfg: &VidCfg,
1418) -> Assoc {
1419    let blank = Assoc {
1420        new_inds: Vec::new(),
1421        unmatched: Vec::new(),
1422        empty: Vec::new(),
1423        det_to_matched: HashMap::new(),
1424        trk_to_high_conf: HashMap::new(),
1425    };
1426    if trk_masks.is_empty() {
1427        // no tracks yet: every detection above threshold starts an object
1428        return Assoc { new_inds: (0..det_masks.len()).collect(), ..blank };
1429    }
1430    if det_masks.is_empty() {
1431        let mut unmatched = Vec::new();
1432        let mut empty = Vec::new();
1433        for (i, m) in trk_masks.iter().enumerate() {
1434            if m.iter().any(|&v| v > 0.0) {
1435                unmatched.push(trk_ids[i]);
1436            } else {
1437                empty.push(trk_ids[i]);
1438            }
1439        }
1440        return Assoc { unmatched, empty, ..blank };
1441    }
1442    let dbin: Vec<Vec<u8>> = det_masks
1443        .iter()
1444        .map(|m| m.iter().map(|&v| u8::from(v > 0.0)).collect())
1445        .collect();
1446    let tbin: Vec<Vec<u8>> = trk_masks
1447        .iter()
1448        .map(|m| m.iter().map(|&v| u8::from(v > 0.0)).collect())
1449        .collect();
1450    let mut ious = vec![0f32; det_masks.len() * trk_masks.len()];
1451    for (i, d) in dbin.iter().enumerate() {
1452        for (j, t) in tbin.iter().enumerate() {
1453            ious[i * trk_masks.len() + j] = mask_iou(d, t);
1454        }
1455    }
1456    let at = |i: usize, j: usize| ious[i * trk_masks.len() + j];
1457
1458    let mut unmatched = Vec::new();
1459    let mut empty = Vec::new();
1460    for (j, t) in tbin.iter().enumerate() {
1461        let nonempty = t.iter().any(|&v| v != 0);
1462        let matched = (0..det_masks.len()).any(|i| at(i, j) >= cfg.trk_assoc_iou_thresh);
1463        if !nonempty {
1464            empty.push(trk_ids[j]);
1465        } else if !matched {
1466            unmatched.push(trk_ids[j]);
1467        }
1468    }
1469    let mut new_inds = Vec::new();
1470    let mut det_to_matched = HashMap::new();
1471    let mut trk_to_high_conf = HashMap::new();
1472    for i in 0..det_masks.len() {
1473        let matches_any = (0..trk_masks.len()).any(|j| at(i, j) >= cfg.assoc_iou_thresh);
1474        let is_new = det_scores[i] >= cfg.new_det_thresh && !matches_any;
1475        if is_new {
1476            new_inds.push(i);
1477        }
1478        let ms: Vec<i64> = (0..trk_masks.len())
1479            .filter(|&j| at(i, j) >= cfg.assoc_iou_thresh)
1480            .map(|j| trk_ids[j])
1481            .collect();
1482        let mut best = (0usize, 0f32);
1483        for j in 0..trk_masks.len() {
1484            if at(i, j) > best.1 {
1485                best = (j, at(i, j));
1486            }
1487        }
1488        if det_scores[i] >= cfg.high_conf_thresh && !is_new && best.1 >= cfg.high_iou_thresh {
1489            trk_to_high_conf.insert(trk_ids[best.0], i);
1490        }
1491        det_to_matched.insert(i, ms);
1492    }
1493    Assoc { new_inds, unmatched, empty, det_to_matched, trk_to_high_conf }
1494}
1495
1496// ───────────────────── outputs ─────────────────────
1497
1498impl Sam3Video {
1499    #[allow(clippy::too_many_arguments)]
1500    fn build_outputs(
1501        &self,
1502        s: &Session,
1503        frame_idx: usize,
1504        pre_ids: &[i64],
1505        dets: &[Det],
1506        plan: &Plan,
1507        frame_size: (usize, usize),
1508    ) -> Result<Vec<TrackObs>> {
1509        let (fw, fh) = frame_size;
1510        let mut out: Vec<TrackObs> = Vec::new();
1511        // tracker masks of the objects that existed before this frame
1512        for &id in pre_ids {
1513            if plan.removed.contains(&id) {
1514                continue;
1515            }
1516            let Some(st) = s.objs.get(&id) else { continue };
1517            let Some(o) = st.cond.get(&frame_idx).or_else(|| st.non_cond.get(&frame_idx)) else {
1518                if std::env::var("SAM3_VID_DEBUG").is_ok() {
1519                    eprintln!("[vid] f{frame_idx} out: obj{id} has NO output at this frame");
1520                }
1521                continue;
1522            };
1523            if o.pred_masks.len() != MASK_SIDE * MASK_SIDE {
1524                if std::env::var("SAM3_VID_DEBUG").is_ok() {
1525                    eprintln!("[vid] f{frame_idx} out: obj{id} mask len {} != {}", o.pred_masks.len(), MASK_SIDE * MASK_SIDE);
1526                }
1527                continue;
1528            }
1529            if o.pred_masks.iter().all(|&v| v <= 0.0) {
1530                if std::env::var("SAM3_VID_DEBUG").is_ok() {
1531                    eprintln!("[vid] f{frame_idx} out: obj{id} mask is ALL-NEGATIVE (min {:.3}) -> no observation",
1532                        o.pred_masks.iter().cloned().fold(f32::INFINITY, f32::min));
1533                }
1534                continue;
1535            }
1536            let score = s.score.get(&id).copied().unwrap_or(0.0);
1537            push_obs(&mut out, id, score, &o.pred_masks, fw, fh, false, true);
1538        }
1539        // masks of objects created on this frame come from the detector
1540        for (k, &di) in plan.new_inds.iter().enumerate() {
1541            let id = plan.new_ids[k];
1542            if plan.removed.contains(&id) {
1543                continue;
1544            }
1545            let mut m = dets[di].mask.clone();
1546            fill_holes_one(&mut m, self.cfg.fill_hole_area);
1547            let score = s.score.get(&id).copied().unwrap_or(dets[di].score);
1548            push_obs(&mut out, id, score, &m, fw, fh, true, false);
1549        }
1550        // hidden objects never reach the outputs
1551        let hidden: HashSet<i64> = s.removed.union(&s.suppressed).copied().collect();
1552        out.retain(|o| !hidden.contains(&o.id));
1553        if out.len() > 1 {
1554            apply_non_overlap(&mut out);
1555        }
1556        out.sort_by_key(|o| o.id);
1557        Ok(out)
1558    }
1559}
1560
1561fn push_obs(
1562    out: &mut Vec<TrackObs>,
1563    id: i64,
1564    score: f32,
1565    logits: &[f32],
1566    fw: usize,
1567    fh: usize,
1568    from_detection: bool,
1569    from_tracker: bool,
1570) {
1571    if logits.len() != MASK_SIDE * MASK_SIDE {
1572        eprintln!("[vid] push_obs: id {id} logits len {} != {MASK_SIDE}", logits.len());
1573        return;
1574    }
1575    let flat = resize2d(logits, MASK_SIDE, MASK_SIDE, fh, fw);
1576    let mask: Vec<u8> = flat.iter().map(|&v| u8::from(v > 0.0)).collect();
1577    if mask.iter().all(|&v| v == 0) {
1578        eprintln!("[vid] push_obs: id {id} empty after resize to {fw}x{fh}");
1579        return;
1580    }
1581    let mut x0 = usize::MAX;
1582    let mut y0 = usize::MAX;
1583    let mut x1 = 0usize;
1584    let mut y1 = 0usize;
1585    for y in 0..fh {
1586        for x in 0..fw {
1587            if mask[y * fw + x] != 0 {
1588                x0 = x0.min(x);
1589                y0 = y0.min(y);
1590                x1 = x1.max(x);
1591                y1 = y1.max(y);
1592            }
1593        }
1594    }
1595    out.push(TrackObs {
1596        id,
1597        score,
1598        box_px: [x0 as f32, y0 as f32, x1 as f32, y1 as f32],
1599        mask,
1600        from_detection,
1601        from_tracker,
1602    });
1603}
1604
1605/// Pixel-wise "highest score wins" between objects of one prompt group
1606/// (`Sam3VideoProcessor._apply_object_wise_non_overlapping_constraints`).
1607fn apply_non_overlap(out: &mut [TrackObs]) {
1608    let len = match out.first() {
1609        Some(o) => o.mask.len(),
1610        None => return,
1611    };
1612    for i in 0..len {
1613        let mut best = usize::MAX;
1614        let mut best_score = f32::NEG_INFINITY;
1615        for (k, o) in out.iter().enumerate() {
1616            if o.mask[i] == 0 {
1617                continue;
1618            }
1619            if o.score > best_score {
1620                best_score = o.score;
1621                best = k;
1622            }
1623        }
1624        if best == usize::MAX {
1625            continue;
1626        }
1627        for (k, o) in out.iter_mut().enumerate() {
1628            if k != best && o.mask[i] != 0 {
1629                o.mask[i] = 0;
1630            }
1631        }
1632    }
1633}
1634
1635// ───────────────────── mask helpers ─────────────────────
1636
1637/// IoU of two binary masks (empty ∪ empty = 1, as in the reference).
1638pub fn mask_iou(a: &[u8], b: &[u8]) -> f32 {
1639    debug_assert_eq!(a.len(), b.len());
1640    let mut inter = 0u32;
1641    let mut union = 0u32;
1642    for (&x, &y) in a.iter().zip(b) {
1643        let x = x != 0;
1644        let y = y != 0;
1645        inter += u32::from(x && y);
1646        union += u32::from(x || y);
1647    }
1648    if union == 0 {
1649        1.0
1650    } else {
1651        inter as f32 / union as f32
1652    }
1653}
1654
1655/// `fill_holes_in_mask_scores`: small background components become 0.1,
1656/// small foreground components (≤ max_area and ≤ half the mask) become -0.1.
1657fn fill_holes(masks: &mut [Vec<f32>], max_area: usize) {
1658    for m in masks.iter_mut() {
1659        fill_holes_one(m, max_area);
1660    }
1661}
1662
1663fn fill_holes_one(mask: &mut [f32], max_area: usize) {
1664    if max_area == 0 || mask.len() != MASK_SIDE * MASK_SIDE {
1665        return;
1666    }
1667    let (w, h) = (MASK_SIDE, MASK_SIDE);
1668    let bg: Vec<bool> = mask.iter().map(|&v| v <= 0.0).collect();
1669    let areas = cc_areas(&bg, w, h);
1670    for i in 0..mask.len() {
1671        if bg[i] && areas[i] > 0 && areas[i] <= max_area {
1672            mask[i] = 0.1;
1673        }
1674    }
1675    let fg: Vec<bool> = mask.iter().map(|&v| v > 0.0).collect();
1676    let total: usize = fg.iter().filter(|&&v| v).count();
1677    let thr = (total / 2).min(max_area);
1678    let areas = cc_areas(&fg, w, h);
1679    for i in 0..mask.len() {
1680        if fg[i] && areas[i] > 0 && areas[i] <= thr {
1681            mask[i] = -0.1;
1682        }
1683    }
1684}
1685
1686/// 4-connected component sizes (every pixel carries its component's area).
1687fn cc_areas(mask: &[bool], w: usize, h: usize) -> Vec<usize> {
1688    let mut label = vec![usize::MAX; mask.len()];
1689    let mut sizes: Vec<usize> = Vec::new();
1690    let mut stack: Vec<usize> = Vec::new();
1691    for start in 0..mask.len() {
1692        if !mask[start] || label[start] != usize::MAX {
1693            continue;
1694        }
1695        let id = sizes.len();
1696        stack.clear();
1697        stack.push(start);
1698        label[start] = id;
1699        let mut count = 0usize;
1700        while let Some(p) = stack.pop() {
1701            count += 1;
1702            let (x, y) = (p % w, p / w);
1703            for (nx, ny) in
1704                [(x.wrapping_sub(1), y), (x + 1, y), (x, y.wrapping_sub(1)), (x, y + 1)]
1705            {
1706                if nx >= w || ny >= h {
1707                    continue;
1708                }
1709                let q = ny * w + nx;
1710                if mask[q] && label[q] == usize::MAX {
1711                    label[q] = id;
1712                    stack.push(q);
1713                }
1714            }
1715        }
1716        sizes.push(count);
1717    }
1718    label
1719        .into_iter()
1720        .map(|l| if l == usize::MAX { 0 } else { sizes[l] })
1721        .collect()
1722}
1723
1724/// `_suppress_object_pw_area_shrinkage`: pixel-wise non-overlap first, then
1725/// drop the objects whose foreground area shrinks below 30 %.
1726fn suppress_area_shrink(masks: &[Vec<f32>]) -> Vec<Vec<f32>> {
1727    if masks.len() <= 1 {
1728        return masks.to_vec();
1729    }
1730    let len = masks[0].len();
1731    let mut non_overlap: Vec<Vec<f32>> = masks.to_vec();
1732    for i in 0..len {
1733        let mut best = 0usize;
1734        let mut best_v = f32::NEG_INFINITY;
1735        for (k, m) in masks.iter().enumerate() {
1736            let v = if m[i] > 0.0 { m[i] } else { -10.0 };
1737            if v > best_v {
1738                best_v = v;
1739                best = k;
1740            }
1741        }
1742        for (k, m) in non_overlap.iter_mut().enumerate() {
1743            if k != best && m[i] > 0.0 {
1744                m[i] = -10.0;
1745            }
1746        }
1747    }
1748    let mut out = masks.to_vec();
1749    for (k, m) in masks.iter().enumerate() {
1750        let before = m.iter().filter(|&&v| v > 0.0).count();
1751        let after = non_overlap[k].iter().filter(|&&v| v > 0.0).count();
1752        if (after as f32) / (before.max(1) as f32) < 0.3 {
1753            for v in out[k].iter_mut() {
1754                if *v > -10.0 {
1755                    *v = -10.0;
1756                }
1757            }
1758        }
1759    }
1760    out
1761}
1762
1763fn relu_vec(v: &[f32]) -> Vec<f32> {
1764    v.iter().map(|x| x.max(0.0)).collect()
1765}
1766
1767fn argmax(v: &[f32]) -> usize {
1768    v.iter()
1769        .enumerate()
1770        .fold((0usize, f32::NEG_INFINITY), |(bi, bv), (i, &x)| {
1771            if x > bv { (i, x) } else { (bi, bv) }
1772        })
1773        .0
1774}
1775
1776/// Depthwise conv (`w [c, 1, k, k]`, `groups = c`, padding k/2).
1777fn depthwise(
1778    w: &[f32],
1779    b: &[f32],
1780    x: &[f32],
1781    c: usize,
1782    h: usize,
1783    wid: usize,
1784    k: usize,
1785    pool: Option<&Pool>,
1786) -> Vec<f32> {
1787    let pad = k / 2;
1788    let mut out = vec![0f32; c * h * wid];
1789    let ob = out.as_mut_ptr() as usize;
1790    let run = |s: usize, e: usize| {
1791        let o = unsafe { std::slice::from_raw_parts_mut(ob as *mut f32, c * h * wid) };
1792        for idx in s..e {
1793            let ch = idx / (h * wid);
1794            let rem = idx % (h * wid);
1795            let (y, x2) = (rem / wid, rem % wid);
1796            let mut acc = b[ch];
1797            for dy in 0..k {
1798                let iy = y as isize + dy as isize - pad as isize;
1799                if iy < 0 || iy >= h as isize {
1800                    continue;
1801                }
1802                for dx in 0..k {
1803                    let ix = x2 as isize + dx as isize - pad as isize;
1804                    if ix < 0 || ix >= wid as isize {
1805                        continue;
1806                    }
1807                    acc += w[(ch * k + dy) * k + dx]
1808                        * x[ch * h * wid + iy as usize * wid + ix as usize];
1809                }
1810            }
1811            o[idx] = acc;
1812        }
1813    };
1814    match pool {
1815        Some(p) => p.run_rows(c * h * wid, &run),
1816        None => run(0, c * h * wid),
1817    }
1818    out
1819}
1820
1821/// `trk.mask_downsample`: 1008² → 252² (conv k4 s4).
1822fn mask_downsample(net: &Net<'_>, mask: &[f32]) -> Result<Vec<f32>> {
1823    let w = net.v("trk.mask_downsample.weight")?;
1824    let b = net.v("trk.mask_downsample.bias")?;
1825    Ok(sam3::conv2d(&w, &b, mask, 1, IMG, IMG, 4, 4, 4, 0, net.pool))
1826}
1827
1828#[cfg(test)]
1829mod tests {
1830    use super::*;
1831
1832    fn cmf_path() -> Option<std::path::PathBuf> {
1833        [" /home/hermes/models/sam3-q8_2f.cmf"]
1834            .iter()
1835            .map(|s| std::path::PathBuf::from(s.trim()))
1836            .find(|p| p.exists())
1837    }
1838
1839    /// The tracker's shape-sensitive half (prompt encoder, mask decoder,
1840    /// memory encoder, memory attention, object pointers) without paying for
1841    /// a detector forward: fabricate the neck features, condition one frame
1842    /// from a synthetic detection mask, then propagate one frame.
1843    #[test]
1844    fn tracker_memory_smoke() {
1845        let Some(path) = cmf_path() else {
1846            eprintln!("SKIP tracker_memory_smoke: cmf not found");
1847            return;
1848        };
1849        let m = Sam3Video::open(path.to_str().expect("path")).expect("open video model");
1850        let net = m.net();
1851
1852        // deterministic synthetic neck features
1853        let fill = |len: usize, k: f32| -> Vec<f32> {
1854            (0..len)
1855                .map(|i| {
1856                    let h = (i as u32).wrapping_mul(2654435761) >> 16;
1857                    (h % 1000) as f32 / 1000.0 * k
1858                })
1859                .collect()
1860        };
1861        let feats = TrkFeats {
1862            l0: fill(32 * MASK_SIDE * MASK_SIDE, 0.05),
1863            l1: fill(64 * 144 * 144, 0.05),
1864            l2: fill(HID * EMB * EMB, 0.05),
1865        };
1866
1867        // synthetic detection: the square at (48..144)² of the 288² grid
1868        let mut mask = vec![0f32; MASK_SIDE * MASK_SIDE];
1869        for y in 48..144 {
1870            for x in 48..144 {
1871                mask[y * MASK_SIDE + x] = 1.0;
1872            }
1873        }
1874        let mut s = Session::new();
1875        s.obj_ids.push(0);
1876        s.max_obj_id = 0;
1877
1878        let (o, high) = m.use_mask_as_output(&net, &feats, &mask).expect("mask-as-output");
1879        assert_eq!(o.pred_masks.len(), MASK_SIDE * MASK_SIDE, "low-res mask shape");
1880        assert_eq!(o.pointer.len(), HID, "object pointer width");
1881        assert_eq!(high.len(), IMG * IMG, "high-res mask shape");
1882        assert!(o.obj_score > 0.0, "object score {}", o.obj_score);
1883        let area = 96 * 96;
1884        let pos = o.pred_masks.iter().filter(|&&v| v > 0.0).count();
1885        assert!(
1886            (pos as i64 - area as i64).abs() <= area as i64 / 4,
1887            "mask-as-output area {pos} vs input {area}"
1888        );
1889
1890        // conditioning frame + memory encode
1891        m.store(&mut s, 0, 0, true, o.clone());
1892        let pix = feats.l2.clone();
1893        m.encode_into(&net, &mut s, 0, vec![(0, o.clone(), high)], &pix, IMG, IMG, true)
1894            .expect("memory encode");
1895        assert!(s.objs[&0].cond[&0].mem_feat.is_some(), "memory features stored");
1896        assert_eq!(s.objs[&0].cond[&0].mem_feat.as_ref().unwrap().len(), EMB * EMB * MEM);
1897
1898        // propagate one frame: memory attention + decoder
1899        let (o1, high1) = m.propagate(&net, &s, 0, 1, &feats).expect("propagate");
1900        assert_eq!(o1.pred_masks.len(), MASK_SIDE * MASK_SIDE);
1901        assert_eq!(high1.len(), IMG * IMG);
1902        let pos1 = o1.pred_masks.iter().filter(|&&v| v > 0.0).count();
1903        let mn = o1.pred_masks.iter().cloned().fold(f32::INFINITY, f32::min);
1904        let mx = o1.pred_masks.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
1905        eprintln!(
1906            "tracker smoke: mask-as-output {} px, propagated {} px (logit {:.2}..{:.2}), obj_score {:.3} -> {:.3}",
1907            pos,
1908            pos1,
1909            mn,
1910            mx,
1911            o.obj_score,
1912            o1.obj_score
1913        );
1914        // the neck features here are synthetic noise, so the propagated mask
1915        // carries no structure — what this test guards is the SHAPE of every
1916        // stage (prompt encoder, mask decoder, memory encoder, memory
1917        // attention, object pointers) plus the object-score sign flip.
1918        assert!(o1.obj_score.is_finite(), "object score must be finite");
1919        assert!(mn.is_finite() && mx.is_finite(), "mask logits must be finite");
1920        assert!(mx > mn, "the decoder must not emit a constant mask ({mn})");
1921    }
1922
1923    fn f32s(p: &std::path::Path) -> Vec<f32> {
1924        let b = std::fs::read(p).unwrap_or_else(|e| panic!("read {}: {e}", p.display()));
1925        b.chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect()
1926    }
1927
1928    fn fx_dir() -> std::path::PathBuf {
1929        std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/sam3_video")
1930    }
1931
1932    /// Reference-anchored tracker check (no detector forward — the fixture
1933    /// supplies the real neck features, memory bank and object pointer):
1934    /// 1. the mask decoder on the frame-0 memory prompt must reproduce the
1935    ///    reference object pointer;
1936    /// 2. memory conditioning + decoder on frame 1 must reproduce the
1937    ///    reference object-score logit (+10.1, "the object is there").
1938    #[test]
1939    fn tracker_fixture_parity() {
1940        let fx = fx_dir();
1941        let Some(path) = cmf_path() else {
1942            eprintln!("SKIP tracker_fixture_parity: cmf not found");
1943            return;
1944        };
1945        if !fx.join("f1/obj0/non_cond_frame_outputs/object_score_logits.f32").exists() {
1946            eprintln!("SKIP tracker_fixture_parity: fixtures missing at {}", fx.display());
1947            return;
1948        }
1949        let m = Sam3Video::open(path.to_str().unwrap()).expect("open");
1950        let net = m.net();
1951        let fill = |len: usize| vec![0f32; len];
1952        let feats0 = TrkFeats { l0: fill(32 * MASK_SIDE * MASK_SIDE), l1: fill(64 * 144 * 144), l2: fill(HID * EMB * EMB) };
1953        let _ = &feats0;
1954
1955        // ── 1. frame 0: decoder with the detection mask prompt → pointer ──
1956        let pix0 = f32s(&fx.join("f0/neck_l2.f32"));
1957        assert_eq!(pix0.len(), EMB * EMB * HID);
1958        let pix0_chw = pix0_to_chw(&pix0); // the fixture stores (HW, C) rows
1959        let hi0 = f32s(&fx.join("f0/obj0/cond_frame_outputs/high_res_masks.f32"));
1960        let bin: Vec<f32> = hi0.iter().map(|&v| if v > 0.0 { 1.0 } else { 0.0 }).collect();
1961        assert_eq!(bin.len(), IMG * IMG);
1962        let ds = mask_downsample(&net, &bin).expect("mask_downsample");
1963        let ds = crate::sam3vid_net::resize2d(&ds, IMG / 4, IMG / 4, MASK_SIDE, MASK_SIDE);
1964        let dense = net.mask_prompt(&ds).expect("mask prompt");
1965        let sparse = net.empty_sparse().expect("sparse");
1966        let pos = net.image_pos().expect("pos");
1967        let (_m, iou0, sc0, toks) = net
1968            .mask_decoder(&pix0_to_chw(&pix0), &pos, &dense, &sparse, (&feats0.l0, &feats0.l1), true)
1969            .expect("decoder f0");
1970        let n = toks.len() / HID;
1971        let best = argmax(&iou0).min(n - 1);
1972        let ptr = m.object_pointer(&net, &toks[best * HID..(best + 1) * HID], sc0 > 0.0).expect("pointer");
1973        let want = f32s(&fx.join("f0/obj0/cond_frame_outputs/object_pointer.f32"));
1974        let cos = dot(&ptr, &want) / (norm(&ptr) * norm(&want) + 1e-9);
1975        eprintln!("fixture f0 pointer: cos={cos:.4} obj_score={sc0:.3} (ref mask prompt)");
1976        assert!(cos >= 0.99, "frame-0 object pointer cosine {cos:.4} vs reference");
1977
1978        // ── 2. frame 1: memory conditioning + decoder → object score ──
1979        let pix1 = f32s(&fx.join("f1/neck_l2.f32"));
1980        let mem_feat = f32s(&fx.join("f0/obj0/cond_frame_outputs/maskmem_features.f32"));
1981        let mem_pos = f32s(&fx.join("f0/obj0/cond_frame_outputs/maskmem_pos_enc.f32"));
1982        let mem_feat_ref = mem_feat.clone();
1983        let memories = vec![(0usize, mem_feat, mem_pos)];
1984        let pointers = vec![(1usize, want.clone())];
1985        let neck_pos = Net::neck_pos();
1986        let cond = m.memory_cond(&net, &pix1, &neck_pos, &memories, &pointers, 8).expect("memory cond");
1987        let cond_ref_path = fx.join("f1/memory_conditioned.f32");
1988        if cond_ref_path.exists() {
1989            let want_cond = f32s(&cond_ref_path);
1990            let cond_cos = dot(&cond, &want_cond) / (norm(&cond) * norm(&want_cond) + 1e-9);
1991            let cond_md = cond.iter().zip(&want_cond).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max);
1992            eprintln!("fixture f1 memory conditioning: cos={cond_cos:.6} maxdiff={cond_md:.4}");
1993            assert!(cond_cos >= 0.99, "frame-1 memory conditioning cosine {cond_cos:.6}");
1994        }
1995        let dense = net.no_mask_prompt().expect("no mask");
1996        let (_m, iou1, sc1, t1) = net
1997            .mask_decoder(&pix0_to_chw(&cond), &pos, &dense, &sparse, (&feats0.l0, &feats0.l1), true)
1998            .expect("decoder f1");
1999        let want_sc = f32s(&fx.join("f1/obj0/non_cond_frame_outputs/object_score_logits.f32"))[0];
2000        eprintln!(
2001            "fixture f1: obj_score ours={sc1:.3} reference={want_sc:.3} (delta {:.3}), iou={iou1:?}",
2002            (sc1 - want_sc).abs()
2003        );
2004        assert!(
2005            (sc1 - want_sc).abs() < 1.0,
2006            "frame-1 object score {sc1:.3} vs reference {want_sc:.3}"
2007        );
2008        assert!(sc1 > 0.0, "the tracked object must read as present (got {sc1:.3})");
2009        let best1 = argmax(&iou1).min(t1.len() / HID - 1);
2010        let ptr1 = m
2011            .object_pointer(&net, &t1[best1 * HID..(best1 + 1) * HID], sc1 > 0.0)
2012            .expect("pointer f1");
2013        let want_ptr1 = f32s(&fx.join("f1/obj0/non_cond_frame_outputs/object_pointer.f32"));
2014        for (j, tok) in t1.chunks_exact(HID).enumerate() {
2015            let p = m.object_pointer(&net, tok, sc1 > 0.0).expect("pointer candidate f1");
2016            let c = dot(&p, &want_ptr1) / (norm(&p) * norm(&want_ptr1) + 1e-9);
2017            eprintln!("fixture f1 pointer candidate {j}: cos={c:.5} iou={:.6}", iou1[j]);
2018        }
2019        let ptr_cos1 = dot(&ptr1, &want_ptr1) / (norm(&ptr1) * norm(&want_ptr1) + 1e-9);
2020        eprintln!("fixture f1 pointer: cos={ptr_cos1:.5}");
2021
2022        // ── 2b. frame 2: TWO spatial memories must each get the grid RoPE ──
2023        // This catches a cross-attention failure that frame 1 cannot see: HF
2024        // repeats the 72×72 frequencies for every memory frame, while the
2025        // object pointers at the end stay unrotated.
2026        let pix2 = f32s(&fx.join("f2/neck_l2.f32"));
2027        let memories2 = vec![
2028            (0usize, mem_feat_ref.clone(), f32s(&fx.join("f0/obj0/cond_frame_outputs/maskmem_pos_enc.f32"))),
2029            (
2030                1usize,
2031                f32s(&fx.join("f1/obj0/non_cond_frame_outputs/maskmem_features.f32")),
2032                f32s(&fx.join("f1/obj0/non_cond_frame_outputs/maskmem_pos_enc.f32")),
2033            ),
2034        ];
2035        let pointers2 = vec![(2usize, want.clone()), (1usize, want_ptr1.clone())];
2036        let cond2 = m.memory_cond(&net, &pix2, &neck_pos, &memories2, &pointers2, 8).expect("memory cond f2");
2037        let cond2_ref_path = fx.join("f2/memory_conditioned.f32");
2038        if cond2_ref_path.exists() {
2039            let want_cond2 = f32s(&cond2_ref_path);
2040            let cond_cos2 = dot(&cond2, &want_cond2) / (norm(&cond2) * norm(&want_cond2) + 1e-9);
2041            let cond_md2 = cond2.iter().zip(&want_cond2).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max);
2042            eprintln!("fixture f2 memory conditioning: cos={cond_cos2:.6} maxdiff={cond_md2:.4}");
2043            assert!(cond_cos2 >= 0.99, "frame-2 memory conditioning cosine {cond_cos2:.6}");
2044        }
2045        let (_m2, iou2, sc2, t2) = net
2046            .mask_decoder(&pix0_to_chw(&cond2), &pos, &dense, &sparse, (&feats0.l0, &feats0.l1), true)
2047            .expect("decoder f2");
2048        let want_sc2 = f32s(&fx.join("f2/obj0/non_cond_frame_outputs/object_score_logits.f32"))[0];
2049        let best2 = argmax(&iou2).min(t2.len() / HID - 1);
2050        let ptr2 = m
2051            .object_pointer(&net, &t2[best2 * HID..(best2 + 1) * HID], sc2 > 0.0)
2052            .expect("pointer f2");
2053        let want_ptr2 = f32s(&fx.join("f2/obj0/non_cond_frame_outputs/object_pointer.f32"));
2054        let ptr_cos2 = dot(&ptr2, &want_ptr2) / (norm(&ptr2) * norm(&want_ptr2) + 1e-9);
2055        eprintln!("fixture f2: obj_score ours={sc2:.3} reference={want_sc2:.3}, pointer cos={ptr_cos2:.5}");
2056        assert!(ptr_cos1 >= 0.99, "frame-1 object pointer cosine {ptr_cos1:.5}");
2057        assert!((sc2 - want_sc2).abs() < 1.0, "frame-2 object score {sc2:.3} vs {want_sc2:.3}");
2058        assert!(ptr_cos2 >= 0.99, "frame-2 object pointer cosine {ptr_cos2:.5}");
2059
2060        // ── 3. our own memory encoder vs the reference's stored memory ──
2061        // (the live pipeline feeds the tracker OUR encoding of frame 0, so it
2062        // must reproduce `maskmem_features` for the same inputs)
2063        for from_pts in [true, false] {
2064            let (mf, mp) = m
2065                .encode_memory(&net, &pix0_chw, &hi0, IMG, IMG, sc0, from_pts)
2066                .expect("encode_memory");
2067            let cos = dot(&mf, &mem_feat_ref) / (norm(&mf) * norm(&mem_feat_ref) + 1e-9);
2068            let md = mf
2069                .iter()
2070                .zip(&mem_feat_ref)
2071                .map(|(a, b)| (a - b).abs())
2072                .fold(0f32, f32::max);
2073            eprintln!(
2074                "fixture mem encode from_pts={from_pts}: cos={cos:.5} maxdiff={md:.4} rms={:.4} (ref rms {:.4})",
2075                norm(&mf) / (mf.len() as f32).sqrt(),
2076                norm(&mem_feat_ref) / (mem_feat_ref.len() as f32).sqrt()
2077            );
2078            assert!(cos >= 0.99, "memory features cosine {cos:.5} (from_pts={from_pts})");
2079            let want_mp = f32s(&fx.join("f0/obj0/cond_frame_outputs/maskmem_pos_enc.f32"));
2080            let pos_cos = dot(&mp, &want_mp) / (norm(&mp) * norm(&want_mp) + 1e-9);
2081            assert!(pos_cos >= 0.99, "memory position cosine {pos_cos:.5}");
2082        }
2083    }
2084
2085    fn pix0_to_chw(rows: &[f32]) -> Vec<f32> {
2086        sam3::rows_to_chw(rows, EMB, EMB)
2087    }
2088
2089    fn dot(a: &[f32], b: &[f32]) -> f32 {
2090        a.iter().zip(b).map(|(x, y)| x * y).sum()
2091    }
2092
2093    fn norm(a: &[f32]) -> f32 {
2094        a.iter().map(|x| x * x).sum::<f32>().sqrt()
2095    }
2096}