[SAM2] Fix inconsistent results with original implementation with input boxes (#40800)

* Fix inconsistencies with box input inference with original repo

* remove print

* always pad

* fix modular

Author: yonigozlan

src/transformers/models/metaclip_2/modeling_metaclip_2.py

@@ -960,9 +960,8 @@ def forward(
         interpolate_pos_encoding: bool = False,
     ) -> MetaClip2Output:
         r"""
-        Args:
-            return_loss (`bool`, *optional*):
-                Whether or not to return the contrastive loss.
+        return_loss (`bool`, *optional*):
+            Whether or not to return the contrastive loss.
 
         Examples:
 

src/transformers/models/metaclip_2/modular_metaclip_2.py

@@ -551,9 +551,8 @@ def forward(
         interpolate_pos_encoding: bool = False,
     ):
         r"""
-        Args:
-            return_loss (`bool`, *optional*):
-                Whether or not to return the contrastive loss.
+        return_loss (`bool`, *optional*):
+            Whether or not to return the contrastive loss.
 
         Examples:
 

src/transformers/models/sam2/modeling_sam2.py

@@ -793,13 +793,14 @@ def _embed_points(self, points: torch.Tensor, labels: torch.Tensor, pad: bool) -
 
     def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor:
         """Embeds box prompts."""
-        boxes = boxes + 0.5  # Shift to center of pixel
-        batch_size, nb_boxes = boxes.shape[:2]
-        coords = boxes.reshape(batch_size, nb_boxes, 2, 2)
-        input_shape = (self.input_image_size, self.input_image_size)
-        corner_embedding = self.shared_embedding(coords, input_shape)
+        boxes += 0.5  # Shift to center of pixel
+        coords = boxes.view(*boxes.shape[:2], 2, 2)
+        # add padding point for consistency with the original implementation
+        coords = torch.nn.functional.pad(coords, (0, 0, 0, 1), mode="constant", value=0)
+        corner_embedding = self.shared_embedding(coords, (self.input_image_size, self.input_image_size))
         corner_embedding[:, :, 0, :] += self.point_embed.weight[2]
         corner_embedding[:, :, 1, :] += self.point_embed.weight[3]
+        corner_embedding[:, :, 2, :] = self.not_a_point_embed.weight.expand_as(corner_embedding[:, :, 2, :])
         return corner_embedding
 
     def forward(

src/transformers/models/sam2/modular_sam2.py

@@ -882,13 +882,14 @@ def _embed_points(self, points: torch.Tensor, labels: torch.Tensor, pad: bool) -
 
     def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor:
         """Embeds box prompts."""
-        boxes = boxes + 0.5  # Shift to center of pixel
-        batch_size, nb_boxes = boxes.shape[:2]
-        coords = boxes.reshape(batch_size, nb_boxes, 2, 2)
-        input_shape = (self.input_image_size, self.input_image_size)
-        corner_embedding = self.shared_embedding(coords, input_shape)
+        boxes += 0.5  # Shift to center of pixel
+        coords = boxes.view(*boxes.shape[:2], 2, 2)
+        # add padding point for consistency with the original implementation
+        coords = torch.nn.functional.pad(coords, (0, 0, 0, 1), mode="constant", value=0)
+        corner_embedding = self.shared_embedding(coords, (self.input_image_size, self.input_image_size))
         corner_embedding[:, :, 0, :] += self.point_embed.weight[2]
         corner_embedding[:, :, 1, :] += self.point_embed.weight[3]
+        corner_embedding[:, :, 2, :] = self.not_a_point_embed.weight.expand_as(corner_embedding[:, :, 2, :])
         return corner_embedding
 
 

src/transformers/models/sam2_video/modeling_sam2_video.py

@@ -1224,13 +1224,14 @@ def _embed_points(self, points: torch.Tensor, labels: torch.Tensor, pad: bool) -
 
     def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor:
         """Embeds box prompts."""
-        boxes = boxes + 0.5  # Shift to center of pixel
-        batch_size, nb_boxes = boxes.shape[:2]
-        coords = boxes.reshape(batch_size, nb_boxes, 2, 2)
-        input_shape = (self.input_image_size, self.input_image_size)
-        corner_embedding = self.shared_embedding(coords, input_shape)
+        boxes += 0.5  # Shift to center of pixel
+        coords = boxes.view(*boxes.shape[:2], 2, 2)
+        # add padding point for consistency with the original implementation
+        coords = torch.nn.functional.pad(coords, (0, 0, 0, 1), mode="constant", value=0)
+        corner_embedding = self.shared_embedding(coords, (self.input_image_size, self.input_image_size))
         corner_embedding[:, :, 0, :] += self.point_embed.weight[2]
         corner_embedding[:, :, 1, :] += self.point_embed.weight[3]
+        corner_embedding[:, :, 2, :] = self.not_a_point_embed.weight.expand_as(corner_embedding[:, :, 2, :])
         return corner_embedding
 
     def forward(

tests/models/sam2/test_modeling_sam2.py

@@ -901,7 +901,7 @@ def test_inference_batched_images_batched_boxes(self):
         self.assertEqual(outputs.pred_masks.shape, (2, 4, 1, 256, 256))
         torch.testing.assert_close(
             outputs.iou_scores,
-            torch.tensor([[[0.9873], [0.9264], [0.9496], [0.9208]], [[0.9445], [0.9496], [0.9497], [0.9481]]]).to(
+            torch.tensor([[[0.9904], [0.9689], [0.9770], [0.9079]], [[0.9739], [0.9816], [0.9838], [0.9781]]]).to(
                 torch_device
             ),
             atol=1e-4,
@@ -912,16 +912,16 @@ def test_inference_batched_images_batched_boxes(self):
             torch.tensor(
                 [
                     [
-                        [[[-7.6204, -11.9286], [-8.7747, -10.5662]]],
-                        [[[-17.1070, -23.4025], [-20.9608, -19.5600]]],
-                        [[[-20.5766, -29.4410], [-26.0739, -24.3225]]],
-                        [[[-19.7201, -29.0836], [-24.4915, -23.6377]]],
+                        [[[-11.1540, -18.3994], [-12.4230, -17.4403]]],
+                        [[[-19.3144, -29.3947], [-24.6341, -24.1144]]],
+                        [[[-24.2983, -37.6470], [-31.6659, -31.0893]]],
+                        [[[-25.4313, -44.0231], [-34.0903, -34.7447]]],
                     ],
                     [
-                        [[[-18.5259, -23.5202], [-25.1906, -17.2518]]],
-                        [[[-20.1214, -25.4215], [-25.7877, -19.1169]]],
-                        [[[-21.0878, -24.7938], [-27.5625, -19.2650]]],
-                        [[[-20.5210, -22.5343], [-26.0968, -17.7544]]],
+                        [[[-22.5539, -30.4633], [-32.8940, -21.6813]]],
+                        [[[-23.6637, -31.3489], [-32.5095, -22.4442]]],
+                        [[[-25.2987, -30.9999], [-34.6243, -24.1717]]],
+                        [[[-26.3150, -30.5313], [-35.0152, -24.0271]]],
                     ],
                 ]
             ).to(torch_device),