chore: import upstream snapshot with attribution

invokeai/backend/image_util/segment_anything/mask_refinement.py

@@ -0,0 +1,50 @@
+# This file contains utilities for Grounded-SAM mask refinement based on:
+# https://github.com/NielsRogge/Transformers-Tutorials/blob/a39f33ac1557b02ebfb191ea7753e332b5ca933f/Grounding%20DINO/GroundingDINO_with_Segment_Anything.ipynb
+
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+
+
+def mask_to_polygon(mask: npt.NDArray[np.uint8]) -> list[tuple[int, int]]:
+    """Convert a binary mask to a polygon.
+
+    Returns:
+        list[list[int]]: List of (x, y) coordinates representing the vertices of the polygon.
+    """
+    # Find contours in the binary mask.
+    contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+    # Find the contour with the largest area.
+    largest_contour = max(contours, key=cv2.contourArea)
+
+    # Extract the vertices of the contour.
+    polygon = largest_contour.reshape(-1, 2).tolist()
+
+    return polygon
+
+
+def polygon_to_mask(
+    polygon: list[tuple[int, int]], image_shape: tuple[int, int], fill_value: int = 1
+) -> npt.NDArray[np.uint8]:
+    """Convert a polygon to a segmentation mask.
+
+    Args:
+        polygon (list): List of (x, y) coordinates representing the vertices of the polygon.
+        image_shape (tuple): Shape of the image (height, width) for the mask.
+        fill_value (int): Value to fill the polygon with.
+
+    Returns:
+        np.ndarray: Segmentation mask with the polygon filled (with value 255).
+    """
+    # Create an empty mask.
+    mask = np.zeros(image_shape, dtype=np.uint8)
+
+    # Convert polygon to an array of points.
+    pts = np.array(polygon, dtype=np.int32)
+
+    # Fill the polygon with white color (255).
+    cv2.fillPoly(mask, [pts], color=(fill_value,))
+
+    return mask

invokeai/backend/image_util/segment_anything/segment_anything_2_pipeline.py

@@ -0,0 +1,109 @@
+from typing import Optional
+
+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
+
+from invokeai.backend.image_util.segment_anything.shared import SAMInput
+from invokeai.backend.raw_model import RawModel
+
+
+class SegmentAnything2Pipeline(RawModel):
+    """A wrapper class for the transformers SAM2 model and processor that makes it compatible with the model manager."""
+
+    def __init__(self, sam2_model: Sam2Model, sam2_processor: Sam2Processor):
+        """Initialize the SAM2 pipeline.
+
+        Args:
+            sam2_model: The SAM2 model
+            sam2_processor: The SAM2 processor (can be Sam2Processor or Sam2VideoProcessor)
+        """
+        self._sam2_model = sam2_model
+        self._sam2_processor = sam2_processor
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
+        # HACK: The SAM2 pipeline may not work on MPS devices. We only allow it to be moved to CPU or CUDA.
+        if device is not None and device.type not in {"cpu", "cuda"}:
+            device = None
+        self._sam2_model.to(device=device, dtype=dtype)
+
+    def calc_size(self) -> int:
+        # HACK: Fix the circular import issue.
+        from invokeai.backend.model_manager.load.model_util import calc_module_size
+
+        return calc_module_size(self._sam2_model)
+
+    def segment(
+        self,
+        image: Image.Image,
+        inputs: list[SAMInput],
+    ) -> torch.Tensor:
+        """Segment the image using the provided inputs.
+
+        Args:
+            image: The image to segment.
+            inputs: A list of SAMInput objects containing bounding boxes and/or point lists.
+
+        Returns:
+            torch.Tensor: The segmentation masks. dtype: torch.bool. shape: [num_masks, channels, height, width].
+        """
+
+        input_boxes: list[list[float]] = []
+        input_points: list[list[list[float]]] = []
+        input_labels: list[list[int]] = []
+
+        for i in inputs:
+            box: list[float] | None = None
+            points: list[list[float]] | None = None
+            labels: list[int] | None = None
+
+            if i.bounding_box is not None:
+                box: list[float] | None = [
+                    i.bounding_box.x_min,
+                    i.bounding_box.y_min,
+                    i.bounding_box.x_max,
+                    i.bounding_box.y_max,
+                ]
+
+            if i.points is not None:
+                points = []
+                labels = []
+                for point in i.points:
+                    points.append([point.x, point.y])
+                    labels.append(point.label.value)
+
+            if box is not None:
+                input_boxes.append(box)
+            if points is not None:
+                input_points.append(points)
+            if labels is not None:
+                input_labels.append(labels)
+
+        batched_input_boxes = [input_boxes] if input_boxes else None
+        batched_input_points = [input_points] if input_points else None
+        batched_input_labels = [input_labels] if input_labels else None
+
+        processed_inputs = self._sam2_processor(
+            images=image,
+            input_boxes=batched_input_boxes,
+            input_points=batched_input_points,
+            input_labels=batched_input_labels,
+            return_tensors="pt",
+        ).to(self._sam2_model.device)
+
+        # Generate masks using the SAM2 model
+        outputs = self._sam2_model(**processed_inputs)
+
+        # 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,
+        )
+
+        # There should be only one batch.
+        assert len(masks) == 1
+        return masks[0]

invokeai/backend/image_util/segment_anything/segment_anything_pipeline.py

@@ -0,0 +1,97 @@
+from typing import Optional
+
+import torch
+from PIL import Image
+from transformers.models.sam import SamModel
+from transformers.models.sam.processing_sam import SamProcessor
+
+from invokeai.backend.image_util.segment_anything.shared import SAMInput
+from invokeai.backend.raw_model import RawModel
+
+
+class SegmentAnythingPipeline(RawModel):
+    """A wrapper class for the transformers SAM model and processor that makes it compatible with the model manager."""
+
+    def __init__(self, sam_model: SamModel, sam_processor: SamProcessor):
+        self._sam_model = sam_model
+        self._sam_processor = sam_processor
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
+        # HACK(ryand): The SAM pipeline does not work on MPS devices. We only allow it to be moved to CPU or CUDA.
+        if device is not None and device.type not in {"cpu", "cuda"}:
+            device = None
+        self._sam_model.to(device=device, dtype=dtype)
+
+    def calc_size(self) -> int:
+        # HACK(ryand): Fix the circular import issue.
+        from invokeai.backend.model_manager.load.model_util import calc_module_size
+
+        return calc_module_size(self._sam_model)
+
+    def segment(
+        self,
+        image: Image.Image,
+        inputs: list[SAMInput],
+    ) -> torch.Tensor:
+        """Segment the image using the provided inputs.
+
+        Args:
+            image: The image to segment.
+            inputs: A list of SAMInput objects containing bounding boxes and/or point lists.
+
+        Returns:
+            torch.Tensor: The segmentation masks. dtype: torch.bool. shape: [num_masks, channels, height, width].
+        """
+
+        input_boxes: list[list[float]] = []
+        input_points: list[list[list[float]]] = []
+        input_labels: list[list[int]] = []
+
+        for i in inputs:
+            box: list[float] | None = None
+            points: list[list[float]] | None = None
+            labels: list[int] | None = None
+
+            if i.bounding_box is not None:
+                box: list[float] | None = [
+                    i.bounding_box.x_min,
+                    i.bounding_box.y_min,
+                    i.bounding_box.x_max,
+                    i.bounding_box.y_max,
+                ]
+
+            if i.points is not None:
+                points = []
+                labels = []
+                for point in i.points:
+                    points.append([point.x, point.y])
+                    labels.append(point.label.value)
+
+            if box is not None:
+                input_boxes.append(box)
+            if points is not None:
+                input_points.append(points)
+            if labels is not None:
+                input_labels.append(labels)
+
+        batched_input_boxes = [input_boxes] if input_boxes else None
+        batched_input_points = input_points if input_points else None
+        batched_input_labels = input_labels if input_labels else None
+
+        processed_inputs = self._sam_processor(
+            images=image,
+            input_boxes=batched_input_boxes,
+            input_points=batched_input_points,
+            input_labels=batched_input_labels,
+            return_tensors="pt",
+        ).to(self._sam_model.device)
+        outputs = self._sam_model(**processed_inputs)
+        masks = self._sam_processor.post_process_masks(
+            masks=outputs.pred_masks,
+            original_sizes=processed_inputs.original_sizes,
+            reshaped_input_sizes=processed_inputs.reshaped_input_sizes,
+        )
+
+        # There should be only one batch.
+        assert len(masks) == 1
+        return masks[0]

invokeai/backend/image_util/segment_anything/shared.py

@@ -0,0 +1,49 @@
+from enum import Enum
+
+from pydantic import BaseModel, model_validator
+from pydantic.fields import Field
+
+
+class BoundingBox(BaseModel):
+    x_min: int = Field(..., description="The minimum x-coordinate of the bounding box (inclusive).")
+    x_max: int = Field(..., description="The maximum x-coordinate of the bounding box (exclusive).")
+    y_min: int = Field(..., description="The minimum y-coordinate of the bounding box (inclusive).")
+    y_max: int = Field(..., description="The maximum y-coordinate of the bounding box (exclusive).")
+
+    @model_validator(mode="after")
+    def check_coords(self):
+        if self.x_min > self.x_max:
+            raise ValueError(f"x_min ({self.x_min}) is greater than x_max ({self.x_max}).")
+        if self.y_min > self.y_max:
+            raise ValueError(f"y_min ({self.y_min}) is greater than y_max ({self.y_max}).")
+        return self
+
+    def tuple(self) -> tuple[int, int, int, int]:
+        """
+        Returns the bounding box as a tuple suitable for use with PIL's `Image.crop()` method.
+        This method returns a tuple of the form (left, upper, right, lower) == (x_min, y_min, x_max, y_max).
+        """
+        return (self.x_min, self.y_min, self.x_max, self.y_max)
+
+
+class SAMPointLabel(Enum):
+    negative = -1
+    neutral = 0
+    positive = 1
+
+
+class SAMPoint(BaseModel):
+    x: int = Field(..., description="The x-coordinate of the point")
+    y: int = Field(..., description="The y-coordinate of the point")
+    label: SAMPointLabel = Field(..., description="The label of the point")
+
+
+class SAMInput(BaseModel):
+    bounding_box: BoundingBox | None = Field(None, description="The bounding box to use for segmentation")
+    points: list[SAMPoint] | None = Field(None, description="The points to use for segmentation")
+
+    @model_validator(mode="after")
+    def check_input(self):
+        if not self.bounding_box and not self.points:
+            raise ValueError("Either bounding_box or points must be provided")
+        return self