slice 14: Object-Multiplex batching — track N objects per frame (faithful) (#21)

Extend Sam3VideoSession to seed and track multiple objects together as the batch
dimension B (Object Multiplex): each object carries its own memory bank + box prompt;
per frame they run through one track_step over the shared, broadcast frame features.
Reuses the slice-12 track_step batch axis rather than retrofitting the old lineage's
bucket mux/demux — identical per-object semantics, simpler path.

Injection-based multi-object tests (CI-safe per the mlx-cpu rule): N objects batch
through one pass, stable ordered ids, per-object mask rows, determinism, and a guard
rejecting mixed seed frames (per-object seed frames / mid-clip spawning is the
association slice). Full regression 608 passed / 13 skipped.

Author: benjipeng

src/mlx_cv/models/sam3/real_video_streaming.py

@@ -31,12 +31,13 @@ class Sam3VideoFrameResult:
     object_score_logits: mx.array  # [num_obj]
 
 
-def _box_to_corner_points(box: tuple[float, float, float, float]) -> tuple[mx.array, mx.array]:
-    """xyxy box (prompt-encoder input-image coords) -> SAM corner points (TL=2, BR=3)."""
+def _box_corner_points(box: tuple[float, float, float, float]) -> mx.array:
+    """xyxy box (prompt-encoder input-image coords) -> 2 SAM corner points ``[2, 2]``."""
     x0, y0, x1, y1 = (float(v) for v in box)
-    coords = mx.array([[[x0, y0], [x1, y1]]])  # [1, 2, 2]
-    labels = mx.array([[2.0, 3.0]])  # [1, 2]
-    return coords, labels
+    return mx.array([[x0, y0], [x1, y1]])
+
+
+_BOX_CORNER_LABELS = (2.0, 3.0)  # SAM box corners: top-left=2, bottom-right=3
 
 
 @dataclass
@@ -53,7 +54,7 @@ def __init__(self, model: Sam3VideoModel, *, num_maskmem: int | None = None):
         self.model = model
         self.tracker = model.tracker_model
         self.num_maskmem = num_maskmem if num_maskmem is not None else self.tracker.num_maskmem
-        self._prompt: _Prompt | None = None
+        self._prompts: list[_Prompt] = []
 
     @classmethod
     def from_tracker(cls, tracker: Sam3TrackerVideoModel, *, num_maskmem: int | None = None) -> "Sam3VideoSession":
@@ -62,11 +63,12 @@ def from_tracker(cls, tracker: Sam3TrackerVideoModel, *, num_maskmem: int | None
         session.model = None
         session.tracker = tracker
         session.num_maskmem = num_maskmem if num_maskmem is not None else tracker.num_maskmem
-        session._prompt = None
+        session._prompts = []
         return session
 
     def add_box_prompt(self, frame_index: int, box, object_id: int) -> None:
-        self._prompt = _Prompt(frame_index=int(frame_index), box=tuple(box), object_id=int(object_id))
+        """Seed an object with a box prompt. Multiple objects may be added (multiplex batch)."""
+        self._prompts.append(_Prompt(frame_index=int(frame_index), box=tuple(box), object_id=int(object_id)))
 
     def propagate(self, pixel_values_per_frame) -> list[Sam3VideoFrameResult]:
         """Extract per-frame tracker features via the detector, then run the streaming loop."""
@@ -75,34 +77,60 @@ def propagate(self, pixel_values_per_frame) -> list[Sam3VideoFrameResult]:
         features = [self.model.extract_tracker_features(pv) for pv in pixel_values_per_frame]
         return self.run_from_features(features)
 
+    @staticmethod
+    def _to_object_batch(x: mx.array, num_objects: int) -> mx.array:
+        """Broadcast a single shared-frame tensor ``[1, ...]`` to the object batch ``[B, ...]``."""
+        if x.shape[0] == num_objects:
+            return x
+        if x.shape[0] != 1:
+            raise ValueError(f"expected frame tensor with leading dim 1 or {num_objects}, got {x.shape[0]}")
+        return mx.repeat(x, num_objects, axis=0)
+
+    def _seed_point_inputs(self, prompts: list[_Prompt]) -> tuple[mx.array, mx.array]:
+        """Batched box-corner prompts for the seed frame -> ``(coords [B,2,2], labels [B,2])``."""
+        coords = mx.stack([_box_corner_points(p.box) for p in prompts], axis=0)
+        labels = mx.broadcast_to(mx.array([_BOX_CORNER_LABELS]), (len(prompts), len(_BOX_CORNER_LABELS)))
+        return coords, labels
+
     def run_from_features(self, per_frame_features) -> list[Sam3VideoFrameResult]:
         """Run the memory-propagation loop over pre-extracted per-frame features.
 
-        Each entry is ``(image_embeddings, image_positional_embeddings, high_res_features)``;
-        ``high_res_features`` is ``[4g-res, 2g-res]`` raw FPN levels (NHWC).
+        Each entry is ``(image_embeddings, image_positional_embeddings, high_res_features)`` for a
+        single frame (shared across objects); ``high_res_features`` is ``[4g-res, 2g-res]`` raw FPN
+        levels (NHWC). All seeded objects are tracked together as the batch dimension ``B`` (Object
+        Multiplex): each object carries its own memory + box prompt; per frame they run through one
+        ``track_step`` over the shared frame features.
         """
-        if self._prompt is None:
+        if not self._prompts:
             raise ValueError("Sam3VideoSession requires a prompt; call add_box_prompt first")
-        prompt = self._prompt
+        seed_frames = {p.frame_index for p in self._prompts}
+        if len(seed_frames) != 1:
+            raise ValueError(
+                "Sam3VideoSession batches objects that share one seed frame; per-object seed frames / "
+                "mid-clip spawning is handled by the association layer (later slice)"
+            )
+        prompts = self._prompts
+        seed_frame = prompts[0].frame_index
+        object_ids = [p.object_id for p in prompts]
+        num_objects = len(prompts)
         bank: list[Sam3TrackerStageOutput] = []
         results: list[Sam3VideoFrameResult] = []
         for frame_index, (vision_features, vision_pos, high_res_features) in enumerate(per_frame_features):
-            is_init = frame_index == prompt.frame_index
-            point_inputs = _box_to_corner_points(prompt.box) if is_init else None
+            is_init = frame_index == seed_frame
             out = self.tracker.track_step(
-                vision_features=vision_features,
-                vision_pos=vision_pos,
-                high_res_features=high_res_features,
+                vision_features=self._to_object_batch(vision_features, num_objects),
+                vision_pos=self._to_object_batch(vision_pos, num_objects),
+                high_res_features=[self._to_object_batch(h, num_objects) for h in high_res_features],
                 is_init_cond_frame=is_init,
-                point_inputs=point_inputs,
+                point_inputs=self._seed_point_inputs(prompts) if is_init else None,
                 previous_frames=bank if bank else None,
             )
             bank.append(out)
             results.append(
                 Sam3VideoFrameResult(
                     frame_index=frame_index,
-                    object_ids=[prompt.object_id],
-                    masks=out.low_res_masks[:, :, :, 0] > 0,
+                    object_ids=list(object_ids),
+                    masks=out.low_res_masks[:, :, :, 0] > 0,  # [B, h, w] (one row per object)
                     object_score_logits=out.object_score_logits.reshape(-1),
                 )
             )

tests/test_sam3_video_multiplex_real.py

@@ -0,0 +1,87 @@
+"""Slice 14: faithful SAM3 Object-Multiplex batching (multiple objects per frame).
+
+Weight-free structural verification that several seeded objects track together as the
+batch dimension ``B`` through one ``track_step`` per frame — each object carries its own
+memory bank + box prompt over the shared frame features. Injection-based (no detector),
+so it runs on the Linux mlx[cpu] CI; numeric parity is the deferred out-of-sandbox gate.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import pytest
+
+from mlx_cv.models.sam3.real_tracker_decoder import Sam3TrackerPromptEncoderConfig
+from mlx_cv.models.sam3.real_video_config import Sam3TrackerVideoConfig
+from mlx_cv.models.sam3.real_video_model import Sam3TrackerVideoModel
+from mlx_cv.models.sam3.real_video_streaming import Sam3VideoSession
+
+CHANNELS = 256
+
+
+def _tiny_tracker(g: int = 4) -> Sam3TrackerVideoModel:
+    config = Sam3TrackerVideoConfig(
+        prompt_encoder=Sam3TrackerPromptEncoderConfig(image_size=g * 16, patch_size=16),
+        memory_attention_rope_feat_sizes=(g, g),
+    )
+    tracker = Sam3TrackerVideoModel(config)
+    mx.eval(tracker.parameters())
+    return tracker
+
+
+def _inject_features(g: int = 4, seed: int = 0):
+    keys = mx.random.split(mx.random.key(seed), 4)
+    return (
+        mx.random.normal((1, g, g, CHANNELS), key=keys[0]),
+        mx.random.normal((1, g, g, CHANNELS), key=keys[1]),
+        [
+            mx.random.normal((1, g * 4, g * 4, CHANNELS), key=keys[2]),
+            mx.random.normal((1, g * 2, g * 2, CHANNELS), key=keys[3]),
+        ],
+    )
+
+
+def test_multiplex_two_objects_batch_through_one_track_step():
+    g = 4
+    session = Sam3VideoSession.from_tracker(_tiny_tracker(g))
+    session.add_box_prompt(0, [4, 4, 28, 28], object_id=5)
+    session.add_box_prompt(0, [20, 20, 60, 60], object_id=9)
+    results = session.run_from_features([_inject_features(g, seed=i) for i in range(3)])
+    assert [r.frame_index for r in results] == [0, 1, 2]
+    for r in results:
+        assert r.object_ids == [5, 9]  # stable, ordered
+        assert r.masks.shape == (2, g * 4, g * 4)  # one mask row per object
+        assert r.object_score_logits.shape == (2,)
+        assert bool(mx.all(mx.isfinite(r.object_score_logits)).item())
+
+
+def test_multiplex_scales_to_more_objects():
+    g = 4
+    session = Sam3VideoSession.from_tracker(_tiny_tracker(g))
+    for i, oid in enumerate((3, 7, 11, 4)):
+        session.add_box_prompt(0, [2 + i, 2 + i, 30 + i, 30 + i], object_id=oid)
+    results = session.run_from_features([_inject_features(g, seed=i) for i in range(2)])
+    for r in results:
+        assert r.object_ids == [3, 7, 11, 4]
+        assert r.masks.shape == (4, g * 4, g * 4)
+
+
+def test_multiplex_rejects_mixed_seed_frames():
+    session = Sam3VideoSession.from_tracker(_tiny_tracker())
+    session.add_box_prompt(0, [4, 4, 28, 28], object_id=1)
+    session.add_box_prompt(1, [4, 4, 28, 28], object_id=2)
+    with pytest.raises(ValueError, match="seed frame"):
+        session.run_from_features([_inject_features() for _ in range(2)])
+
+
+def test_multiplex_is_deterministic():
+    g = 4
+    session = Sam3VideoSession.from_tracker(_tiny_tracker(g))
+    session.add_box_prompt(0, [4, 4, 28, 28], object_id=1)
+    session.add_box_prompt(0, [20, 20, 60, 60], object_id=2)
+    frames = [_inject_features(g, seed=i) for i in range(3)]
+    first = session.run_from_features(frames)
+    second = session.run_from_features(frames)
+    for a, b in zip(first, second):
+        assert bool(mx.all(a.masks == b.masks).item())
+        assert bool(mx.all(a.object_score_logits == b.object_score_logits).item())