feat(app): update CanvasV2MaskAndCropInvocation

This commit is contained in:
psychedelicious 2024-07-16 21:31:43 +10:00
parent 010916158b
commit f8f21c0edd

View File

@ -1017,8 +1017,8 @@ class MaskFromIDInvocation(BaseInvocation, WithMetadata, WithBoard):
@invocation_output("canvas_v2_mask_and_crop_output") @invocation_output("canvas_v2_mask_and_crop_output")
class CanvasV2MaskAndCropOutput(ImageOutput): class CanvasV2MaskAndCropOutput(ImageOutput):
x: int = OutputField(description="The x coordinate of the image") offset_x: int = OutputField(description="The x offset of the image, after cropping")
y: int = OutputField(description="The y coordinate of the image") offset_y: int = OutputField(description="The y offset of the image, after cropping")
@invocation( @invocation(
@ -1029,32 +1029,33 @@ class CanvasV2MaskAndCropOutput(ImageOutput):
version="1.0.0", version="1.0.0",
) )
class CanvasV2MaskAndCropInvocation(BaseInvocation, WithMetadata, WithBoard): class CanvasV2MaskAndCropInvocation(BaseInvocation, WithMetadata, WithBoard):
"""Apply a mask to an image""" """Handles Canvas V2 image output masking and cropping"""
image: ImageField = InputField(description="The image to apply the mask to") image: ImageField = InputField(description="The image to apply the mask to")
mask: ImageField = InputField(description="The mask to apply") mask: ImageField = InputField(description="The mask to apply")
invert: bool = InputField(default=False, description="Whether or not to invert the mask") mask_blur: int = InputField(default=0, ge=0, description="The amount to blur the mask by")
crop_visible: bool = InputField(default=False, description="Crop the image to the mask")
def _prepare_mask(self, mask: Image.Image) -> Image.Image:
mask_array = numpy.array(mask)
kernel = numpy.ones((self.mask_blur, self.mask_blur), numpy.uint8)
dilated_mask_array = cv2.erode(mask_array, kernel, iterations=3)
dilated_mask = Image.fromarray(dilated_mask_array)
if self.mask_blur > 0:
mask = dilated_mask.filter(ImageFilter.GaussianBlur(self.mask_blur))
return ImageOps.invert(mask.convert("L"))
def invoke(self, context: InvocationContext) -> CanvasV2MaskAndCropOutput: def invoke(self, context: InvocationContext) -> CanvasV2MaskAndCropOutput:
image = context.images.get_pil(self.image.image_name) image = context.images.get_pil(self.image.image_name)
mask = context.images.get_pil(self.mask.image_name) mask = self._prepare_mask(context.images.get_pil(self.mask.image_name))
if self.invert:
mask = ImageOps.invert(mask)
image.putalpha(mask) image.putalpha(mask)
bbox = image.getbbox() bbox = image.getbbox()
image = image.crop(bbox)
if self.crop_visible:
image = image.crop(bbox)
image_dto = context.images.save(image=image) image_dto = context.images.save(image=image)
return CanvasV2MaskAndCropOutput( return CanvasV2MaskAndCropOutput(
image=ImageField(image_name=image_dto.image_name), image=ImageField(image_name=image_dto.image_name),
x=bbox[0], offset_x=bbox[0],
y=bbox[1], offset_y=bbox[1],
width=image.width, width=image.width,
height=image.height, height=image.height,
) )