Fix(SAM2): mask post-processing with Transformers 5.x

invokeai/backend/image_util/segment_anything/segment_anything_2_pipeline.py

@@ -2,8 +2,6 @@
 
 import torch
 from PIL import Image
-
-# Import SAM2 components - these should be available in transformers 4.56.0+
 from transformers.models.sam2 import Sam2Model
 from transformers.models.sam2.processing_sam2 import Sam2Processor
 
@@ -19,7 +17,7 @@ def __init__(self, sam2_model: Sam2Model, sam2_processor: Sam2Processor):
 
         Args:
             sam2_model: The SAM2 model
-            sam2_processor: The SAM2 processor (can be Sam2Processor or Sam2VideoProcessor)
+            sam2_processor: The SAM2 processor used to preprocess inputs and post-process masks
         """
         self._sam2_model = sam2_model
         self._sam2_processor = sam2_processor
@@ -99,9 +97,7 @@ def segment(
 
         # Post-process the masks to get the final segmentation
         masks = self._sam2_processor.post_process_masks(
-            masks=outputs.pred_masks,
-            original_sizes=processed_inputs.original_sizes,
-            reshaped_input_sizes=processed_inputs.reshaped_input_sizes,
+            masks=outputs.pred_masks, original_sizes=processed_inputs.original_sizes
         )
 
         # There should be only one batch.

tests/backend/image_util/segment_anything/test_segment_anything_2_pipeline.py

@@ -0,0 +1,53 @@
+import torch
+from PIL import Image
+from transformers.models.sam2.configuration_sam2 import Sam2Config
+from transformers.models.sam2.image_processing_sam2 import Sam2ImageProcessor
+from transformers.models.sam2.modeling_sam2 import Sam2Model
+from transformers.models.sam2.processing_sam2 import Sam2Processor
+
+from invokeai.backend.image_util.segment_anything.segment_anything_2_pipeline import (
+    SegmentAnything2Pipeline,
+)
+from invokeai.backend.image_util.segment_anything.shared import (
+    SAMInput,
+    SAMPoint,
+    SAMPointLabel,
+)
+
+
+def test_segment_anything_2_pipeline_segment():
+    width = 96
+    height = 64
+
+    image = Image.new("RGB", (width, height))
+
+    model = Sam2Model(Sam2Config())
+    processor = Sam2Processor(
+        image_processor=Sam2ImageProcessor(),
+    )
+
+    pipeline = SegmentAnything2Pipeline(
+        sam2_model=model,
+        sam2_processor=processor,
+    )
+
+    inputs = [
+        SAMInput(
+            points=[
+                SAMPoint(
+                    x=width // 2,
+                    y=height // 2,
+                    label=SAMPointLabel.positive,
+                )
+            ]
+        )
+    ]
+
+    with torch.inference_mode():
+        masks = pipeline.segment(
+            image=image,
+            inputs=inputs,
+        )
+
+    assert masks.dtype == torch.bool
+    assert masks.shape == (1, 3, height, width)

tests/backend/image_util/segment_anything/test_segment_anything_pipeline.py

@@ -0,0 +1,53 @@
+import torch
+from PIL import Image
+from transformers.models.sam.configuration_sam import SamConfig
+from transformers.models.sam.image_processing_sam import SamImageProcessor
+from transformers.models.sam.modeling_sam import SamModel
+from transformers.models.sam.processing_sam import SamProcessor
+
+from invokeai.backend.image_util.segment_anything.segment_anything_pipeline import (
+    SegmentAnythingPipeline,
+)
+from invokeai.backend.image_util.segment_anything.shared import (
+    SAMInput,
+    SAMPoint,
+    SAMPointLabel,
+)
+
+
+def test_segment_anything_pipeline_segment():
+    width = 96
+    height = 64
+
+    image = Image.new("RGB", (width, height))
+
+    model = SamModel(SamConfig())
+    processor = SamProcessor(
+        image_processor=SamImageProcessor(),
+    )
+
+    pipeline = SegmentAnythingPipeline(
+        sam_model=model,
+        sam_processor=processor,
+    )
+
+    inputs = [
+        SAMInput(
+            points=[
+                SAMPoint(
+                    x=width // 2,
+                    y=height // 2,
+                    label=SAMPointLabel.positive,
+                )
+            ]
+        )
+    ]
+
+    with torch.inference_mode():
+        masks = pipeline.segment(
+            image=image,
+            inputs=inputs,
+        )
+
+    assert masks.dtype == torch.bool
+    assert masks.shape == (1, 3, height, width)