1use 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
46pub const MASK_SIDE: usize = 288;
49const IMG: usize = 1008;
51const MEM_SIDE: usize = 1152;
53
54#[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 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 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#[derive(Default, Clone)]
199struct ObjOut {
200 pred_masks: Vec<f32>,
202 obj_score: f32,
204 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
216pub struct Session {
218 pub num_frames: usize,
221 pub obj_ids: Vec<i64>,
223 pub max_obj_id: i64,
224 objs: HashMap<i64, ObjState>,
225 score: HashMap<i64, f32>,
227 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 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 pub fn hidden_ids(&self) -> Vec<i64> {
267 self.removed.union(&self.suppressed).copied().collect()
268 }
269}
270
271#[derive(Clone, Debug)]
275pub struct TrackObs {
276 pub id: i64,
277 pub score: f32,
278 pub box_px: [f32; 4],
280 pub mask: Vec<u8>,
282 pub from_detection: bool,
284 pub from_tracker: bool,
286}
287
288#[derive(Clone, Debug)]
290pub struct FrameOut {
291 pub frame: usize,
292 pub tracks: Vec<TrackObs>,
293 pub detections: usize,
295 pub presence: f32,
296 pub det_ms: f64,
297 pub trk_ms: f64,
298 pub total_ms: f64,
299}
300
301struct Det {
304 score: f32,
305 mask: Vec<f32>,
306}
307
308pub 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 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 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 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 let feats = {
363 let _p = crate::cpuprof::spans::span("vid.feats");
364 net.feats(&out.trunk_feats)?
365 };
366
367 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 return Ok(pix.to_vec());
797 }
798 let tpos = net.v("trk.maskmem_tpos_enc")?; 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 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 let mem_len = mem.len() / MEM;
829
830 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 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; 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 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 #[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], mask_logits: &[f32], 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 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 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 for layer in 0..2 {
1026 let lp = format!("trk.maskmem_backbone.fuser.layers.{layer}.");
1027 let block_in = f.clone(); 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 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); 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); 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 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 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
1140struct 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 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 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 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 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 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 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
1401struct 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
1411fn 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 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
1496impl 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 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 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 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
1605fn 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
1635pub 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
1655fn 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
1686fn 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
1724fn 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
1776fn 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
1821fn 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 #[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 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 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 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 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 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 #[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 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); 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 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 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 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}