diff --git a/invokeai/backend/image_util/segment_anything/segment_anything_2_pipeline.py b/invokeai/backend/image_util/segment_anything/segment_anything_2_pipeline.py index c7b6bf6b39b..845c825e3e8 100644 --- a/invokeai/backend/image_util/segment_anything/segment_anything_2_pipeline.py +++ b/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. diff --git a/tests/backend/image_util/segment_anything/test_segment_anything_2_pipeline.py b/tests/backend/image_util/segment_anything/test_segment_anything_2_pipeline.py new file mode 100644 index 00000000000..9aa4f454dd1 --- /dev/null +++ b/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) diff --git a/tests/backend/image_util/segment_anything/test_segment_anything_pipeline.py b/tests/backend/image_util/segment_anything/test_segment_anything_pipeline.py new file mode 100644 index 00000000000..736e592abfa --- /dev/null +++ b/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)