From fdae8d262ad163994b466fa3e516eb1b00cd17d8 Mon Sep 17 00:00:00 2001 From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com> Date: Fri, 1 Mar 2024 03:41:51 +0530 Subject: [PATCH] fix: seg mask bug fixes - Make sure the output more is correct so other nodes can interact with it - Add option to invert mask if needed --- .../app/invocations/controlnet_image_processors.py | 5 ++++- .../image_util/segment_anything/sam_image_predictor.py | 10 +++++++--- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/invokeai/app/invocations/controlnet_image_processors.py b/invokeai/app/invocations/controlnet_image_processors.py index 35f9f26c2f7..0c9ff934d35 100644 --- a/invokeai/app/invocations/controlnet_image_processors.py +++ b/invokeai/app/invocations/controlnet_image_processors.py @@ -596,9 +596,12 @@ class SegmentAnythingImageProcessorInvocation(ImageProcessorInvocation): x_coordinate: int = InputField(default=0, ge=0, description="X-coordinate of your subject") y_coordinate: int = InputField(default=0, ge=0, description="Y-coordinate of your subject") background: bool = InputField(default=False, description="Object to mask is in the background") + invert: bool = InputField(default=False, description="Invert the generated mask") def run_processor(self, image: Image.Image): sam_predictor = SAMImagePredictor() sam_predictor.load_model(self.model_type) - mask = sam_predictor(image, background=self.background, position=(self.x_coordinate, self.y_coordinate)) + mask = sam_predictor( + image, background=self.background, position=(self.x_coordinate, self.y_coordinate), invert=self.invert + ) return mask diff --git a/invokeai/backend/image_util/segment_anything/sam_image_predictor.py b/invokeai/backend/image_util/segment_anything/sam_image_predictor.py index 7646698bb41..1dec516dc2f 100644 --- a/invokeai/backend/image_util/segment_anything/sam_image_predictor.py +++ b/invokeai/backend/image_util/segment_anything/sam_image_predictor.py @@ -2,7 +2,7 @@ from typing import Literal, Tuple import numpy as np -from PIL import Image +from PIL import Image, ImageOps from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.backend.image_util.segment_anything import SamPredictor, sam_model_registry @@ -72,7 +72,9 @@ def load_model(self, model_type: SEGMENT_ANYTHING_MODEL_TYPES = "small"): ) sam_predictor = sam_predictor = SamPredictor(sam_model) - def __call__(self, image: Image.Image, background: bool = False, position: Tuple[int, int] = (0, 0)) -> Image.Image: + def __call__( + self, image: Image.Image, background: bool = False, position: Tuple[int, int] = (0, 0), invert: bool = False + ) -> Image.Image: global sam_predictor input_image = np.array(image.convert("RGB")) if image.mode != "RGB" else np.array(image) @@ -82,7 +84,9 @@ def __call__(self, image: Image.Image, background: bool = False, position: Tuple if sam_predictor: sam_predictor.set_image(input_image) masks, _, _ = sam_predictor.predict(input_point, input_label) - mask = Image.fromarray(masks[0]) + mask = Image.fromarray(masks[0]).convert("RGB") + if invert: + mask = ImageOps.invert(mask) return mask else: return Image.new("RGB", (image.width, image.height), color="black")