Perf/reduce redundant inference work

sam2/automatic_mask_generator.py

@@ -311,28 +311,42 @@ def _process_batch(
         in_labels = torch.ones(
             in_points.shape[0], dtype=torch.int, device=in_points.device
         )
+        # Predicted IoU is available before full-resolution mask work. Keep the
+        # native order for refinement and optional connected-component cleanup.
+        defer_postprocess = (
+            not self.use_m2m
+            and self.pred_iou_thresh > 0.0
+            and self.min_mask_region_area == 0
+        )
         masks, iou_preds, low_res_masks = self.predictor._predict(
             in_points[:, None, :],
             in_labels[:, None],
             multimask_output=self.multimask_output,
             return_logits=True,
+            postprocess=not defer_postprocess,
         )
 
         # Serialize predictions and store in MaskData
         data = MaskData(
             masks=masks.flatten(0, 1),
             iou_preds=iou_preds.flatten(0, 1),
             points=points.repeat_interleave(masks.shape[1], dim=0),
-            low_res_masks=low_res_masks.flatten(0, 1),
         )
-        del masks
+        if self.use_m2m:
+            data["low_res_masks"] = low_res_masks.flatten(0, 1)
+        del masks, low_res_masks
 
         if not self.use_m2m:
             # Filter by predicted IoU
             if self.pred_iou_thresh > 0.0:
                 keep_mask = data["iou_preds"] > self.pred_iou_thresh
                 data.filter(keep_mask)
 
+            if defer_postprocess:
+                data["masks"] = self.predictor._transforms.postprocess_masks(
+                    data["masks"].unsqueeze(1), im_size
+                ).squeeze(1)
+
             # Calculate and filter by stability score
             data["stability_score"] = calculate_stability_score(
                 data["masks"], self.mask_threshold, self.stability_score_offset
@@ -351,6 +365,7 @@ def _process_batch(
             masks, ious = self.refine_with_m2m(
                 in_points, labels, data["low_res_masks"], self.points_per_batch
             )
+            del data["low_res_masks"]
             data["masks"] = masks.squeeze(1)
             data["iou_preds"] = ious.squeeze(1)
 

sam2/sam2_image_predictor.py

@@ -194,6 +194,12 @@ def predict_batch(
         all_masks = []
         all_ious = []
         all_low_res_masks = []
+        # The dense positional encoding is image-independent within this call.
+        # Keep its lifetime local so model/autocast changes need no cache invalidation.
+        image_pe = None
+        if num_images > 1:
+            with torch.no_grad():
+                image_pe = self.model.sam_prompt_encoder.get_dense_pe()
         for img_idx in range(num_images):
             # Transform input prompts
             point_coords = (
@@ -222,6 +228,7 @@ def predict_batch(
                 multimask_output,
                 return_logits=return_logits,
                 img_idx=img_idx,
+                image_pe=image_pe,
             )
             masks_np = masks.squeeze(0).float().detach().cpu().numpy()
             iou_predictions_np = (
@@ -343,6 +350,8 @@ def _predict(
         multimask_output: bool = True,
         return_logits: bool = False,
         img_idx: int = -1,
+        postprocess: bool = True,
+        image_pe: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """
         Predict masks for the given input prompts, using the currently set image.
@@ -369,9 +378,14 @@ def _predict(
             input prompts, multimask_output=False can give better results.
           return_logits (bool): If true, returns un-thresholded masks logits
             instead of a binary mask.
+          postprocess (bool): If false, return raw, unclamped low-resolution
+            logits as the first output for filtering before spatial expansion.
+          image_pe (torch.Tensor or None): Dense positional encoding computed
+            in the current prediction call; None computes it for this image.
 
         Returns:
-          (torch.Tensor): The output masks in BxCxHxW format, where C is the
+          (torch.Tensor): The output masks (raw low-resolution masks when
+            postprocess=False) in BxCxHxW format, where C is the
             number of masks, and (H, W) is the original image size.
           (torch.Tensor): An array of shape BxC containing the model's
             predictions for the quality of each mask.
@@ -419,18 +433,26 @@ def _predict(
         ]
         low_res_masks, iou_predictions, _, _ = self.model.sam_mask_decoder(
             image_embeddings=self._features["image_embed"][img_idx].unsqueeze(0),
-            image_pe=self.model.sam_prompt_encoder.get_dense_pe(),
+            image_pe=(
+                self.model.sam_prompt_encoder.get_dense_pe()
+                if image_pe is None
+                else image_pe
+            ),
             sparse_prompt_embeddings=sparse_embeddings,
             dense_prompt_embeddings=dense_embeddings,
             multimask_output=multimask_output,
             repeat_image=batched_mode,
             high_res_features=high_res_features,
         )
 
-        # Upscale the masks to the original image resolution
-        masks = self._transforms.postprocess_masks(
-            low_res_masks, self._orig_hw[img_idx]
-        )
+        # AMG can defer postprocessing until after its predicted-IoU filter.
+        # Keep raw logits here: refinement logits below are clamped separately.
+        if postprocess:
+            masks = self._transforms.postprocess_masks(
+                low_res_masks, self._orig_hw[img_idx]
+            )
+        else:
+            masks = low_res_masks
         low_res_masks = torch.clamp(low_res_masks, -32.0, 32.0)
         if not return_logits:
             masks = masks > self.mask_threshold

sam2/utils/amg.py

@@ -119,6 +119,25 @@ def mask_to_rle_pytorch(tensor: torch.Tensor) -> List[Dict[str, Any]]:
     diff = tensor[:, 1:] ^ tensor[:, :-1]
     change_indices = diff.nonzero()
 
+    # Sparse mask boundaries are cheap to serialize on CPU in one batch. This
+    # avoids scanning the entire change-index tensor and synchronizing once per
+    # mask. Keep dense boundaries on the original path: transferring both index
+    # columns can cost more than the per-mask run lengths in that regime.
+    # Each boundary has two int64 indices; cap the transfer table at
+    # 256 KiB per mask on average before using the original dense path.
+    if tensor.is_cuda and b > 1 and change_indices.shape[0] <= 16384 * b:
+        changes = change_indices.cpu().numpy()
+        starts = tensor[:, 0].cpu().tolist()
+        boundaries = np.searchsorted(changes[:, 0], np.arange(b + 1))
+        out = []
+        for i in range(b):
+            cur_idxs = changes[boundaries[i] : boundaries[i + 1], 1] + 1
+            counts = np.diff(np.concatenate(([0], cur_idxs, [h * w]))).tolist()
+            if starts[i]:
+                counts.insert(0, 0)
+            out.append({"size": [h, w], "counts": counts})
+        return out
+
     # Encode run length
     out = []
     for i in range(b):

sam2/utils/misc.py

@@ -89,11 +89,17 @@ def mask_to_box(masks: torch.Tensor):
     return bbox_coords
 
 
-def _load_img_as_tensor(img_path, image_size):
+def _load_img_as_tensor(img_path, image_size, *, use_float32=False):
     img_pil = Image.open(img_path)
     img_np = np.array(img_pil.convert("RGB").resize((image_size, image_size)))
     if img_np.dtype == np.uint8:  # np.uint8 is expected for JPEG images
-        img_np = img_np / 255.0
+        if use_float32:
+            # The eager loader stores FP32 frames. Avoid a temporary FP64 frame
+            # and its later conversion; preserve the async loader's default dtype.
+            img_np = img_np.astype(np.float32)
+            img_np /= 255.0
+        else:
+            img_np = img_np / 255.0
     else:
         raise RuntimeError(f"Unknown image dtype: {img_np.dtype} on {img_path}")
     img = torch.from_numpy(img_np).permute(2, 0, 1)
@@ -266,7 +272,9 @@ def load_video_frames_from_jpg_images(
 
     images = torch.zeros(num_frames, 3, image_size, image_size, dtype=torch.float32)
     for n, img_path in enumerate(tqdm(img_paths, desc="frame loading (JPEG)")):
-        images[n], video_height, video_width = _load_img_as_tensor(img_path, image_size)
+        images[n], video_height, video_width = _load_img_as_tensor(
+            img_path, image_size, use_float32=True
+        )
     if not offload_video_to_cpu:
         images = images.to(compute_device)
         img_mean = img_mean.to(compute_device)

tests/test_amg_filtering.py

@@ -0,0 +1,87 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+import torch.nn.functional as F
+
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+from test_amg_lifetime import make_generator, run_batch
+
+
+class EarlyFilteringTests(unittest.TestCase):
+    def test_existing_iou_selection_preserves_results_and_skips_expansion(self):
+        for threshold, expected_count in [(0.5, 1), (1.0, 0), (0.0, 2)]:
+            with self.subTest(threshold=threshold):
+                # The synthetic cleanup implementation is identical; its flag
+                # selects the original expansion-before-filtering ordering.
+                original, full_work = make_generator(threshold=threshold, cleanup=1)
+                filtered, kept_work = make_generator(threshold=threshold)
+                expected = run_batch(original)
+                actual = run_batch(filtered)
+                self.assertEqual(actual._stats.keys(), expected._stats.keys())
+                for key in actual._stats:
+                    if isinstance(actual[key], torch.Tensor):
+                        torch.testing.assert_close(
+                            actual[key], expected[key], rtol=0, atol=0
+                        )
+                    else:
+                        self.assertEqual(actual[key], expected[key])
+                self.assertEqual(sum(full_work), 2)
+                self.assertEqual(sum(kept_work), expected_count)
+
+    def test_refinement_and_cleanup_keep_original_expansion(self):
+        for refine, cleanup in [(True, 0), (False, 1)]:
+            generator, work = make_generator(refine=refine, cleanup=cleanup)
+            run_batch(generator)
+            self.assertEqual(work[0], 2)
+
+    def test_deferred_logits_are_not_clamped_before_interpolation(self):
+        raw = torch.tensor([[[[-50.0, 50.0], [-1.0, 1.0]]]])
+        prompt = SimpleNamespace(get_dense_pe=lambda: torch.zeros(1, 1, 1))
+
+        class Prompt:
+            def __call__(self, **kwargs):
+                return torch.zeros(1, 1, 1), torch.zeros(1, 1, 1)
+
+            get_dense_pe = staticmethod(prompt.get_dense_pe)
+
+        calls = []
+
+        def post(masks, size):
+            calls.append(masks.shape)
+            return F.interpolate(
+                masks.float(), size, mode="bilinear", align_corners=False
+            )
+
+        predictor = SimpleNamespace(
+            _is_image_set=True,
+            _features={"high_res_feats": [], "image_embed": torch.zeros(1, 1, 1, 1)},
+            _orig_hw=[(4, 4)],
+            mask_threshold=0,
+            _transforms=SimpleNamespace(postprocess_masks=post),
+            model=SimpleNamespace(
+                sam_prompt_encoder=Prompt(),
+                sam_mask_decoder=lambda **kwargs: (raw, torch.ones(1, 1), None, None),
+            ),
+        )
+        full, _, refine = SAM2ImagePredictor._predict(
+            predictor, None, None, return_logits=True
+        )
+        deferred, _, refine2 = SAM2ImagePredictor._predict(
+            predictor, None, None, return_logits=True, postprocess=False
+        )
+        self.assertEqual(len(calls), 1)
+        torch.testing.assert_close(deferred, raw, rtol=0, atol=0)
+        torch.testing.assert_close(refine, raw.clamp(-32, 32), rtol=0, atol=0)
+        torch.testing.assert_close(refine, refine2, rtol=0, atol=0)
+        torch.testing.assert_close(full, post(deferred, (4, 4)), rtol=0, atol=0)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_amg_lifetime.py

@@ -0,0 +1,98 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import numpy as np
+
+import sam2.automatic_mask_generator as amg
+import torch
+import torch.nn.functional as F
+from sam2.utils.amg import MaskData, rle_to_mask
+
+
+def make_generator(refine=False, threshold=0.5, cleanup=0):
+    """Synthetic decoder outputs with interior masks and one rejected score."""
+    raw = torch.full((2, 1, 4, 4), -50.0)
+    raw[0, 0, 1:3, 1:3] = 50.0
+    raw[1, 0, 1:2, 1:3] = 50.0
+    scores = torch.tensor([[0.9], [0.1]])
+    expanded = []
+
+    def postprocess(masks, size):
+        expanded.append(masks.shape[0] * masks.shape[1])
+        return F.interpolate(masks, size, mode="bilinear", align_corners=False)
+
+    def predict(*args, postprocess=True, **kwargs):
+        masks = transforms.postprocess_masks(raw, (16, 16)) if postprocess else raw
+        return masks.clone(), scores.clone(), raw.clamp(-32, 32)
+
+    transforms = SimpleNamespace(
+        transform_coords=lambda points, **kwargs: points,
+        postprocess_masks=postprocess,
+    )
+    generator = object.__new__(amg.SAM2AutomaticMaskGenerator)
+    generator.predictor = SimpleNamespace(
+        device=torch.device("cpu"), _transforms=transforms, _predict=predict
+    )
+    generator.use_m2m = refine
+    generator.pred_iou_thresh = threshold
+    generator.min_mask_region_area = cleanup
+    generator.stability_score_thresh = 0.0
+    generator.stability_score_offset = 1.0
+    generator.mask_threshold = 0.0
+    generator.multimask_output = True
+    generator.points_per_batch = 2
+
+    def refine_masks(points, labels, low_res, batch):
+        torch.testing.assert_close(
+            low_res, raw.flatten(0, 1).clamp(-32, 32), rtol=0, atol=0
+        )
+        return postprocess(raw, (16, 16)), scores.clone()
+
+    generator.refine_with_m2m = refine_masks
+    return generator, expanded
+
+
+def run_batch(generator):
+    return generator._process_batch(
+        np.array([[4, 4], [8, 8]], dtype=np.float32),
+        (16, 16),
+        [0, 0, 16, 16],
+        (16, 16),
+        normalize=False,
+    )
+
+
+class MaskLifetimeTests(unittest.TestCase):
+    def test_downstream_filters_do_not_retain_refinement_logits(self):
+        class CheckedData(MaskData):
+            def filter(self, keep):
+                if "low_res_masks" in self._stats:
+                    raise AssertionError(
+                        "refinement logits reached a downstream filter"
+                    )
+                return super().filter(keep)
+
+        for refine in (False, True):
+            with self.subTest(refine=refine):
+                generator, _ = make_generator(refine=refine)
+                with patch.object(amg, "MaskData", CheckedData):
+                    result = run_batch(generator)
+                self.assertNotIn("low_res_masks", result._stats)
+                self.assertEqual(len(result["rles"]), 1)
+                expected = np.zeros((16, 16), dtype=bool)
+                expected[4:12, 4:12] = True
+                # Bilinear interpolation rounds the four corners inward.
+                expected[4, 4] = expected[4, 11] = False
+                expected[11, 4] = expected[11, 11] = False
+                np.testing.assert_array_equal(rle_to_mask(result["rles"][0]), expected)
+
+
+if __name__ == "__main__":
+    unittest.main()