From d5e824e782376168a4872135e36289139a25f1b6 Mon Sep 17 00:00:00 2001 From: Ryan Dick Date: Thu, 25 Jul 2024 18:17:18 -0400 Subject: [PATCH] Add mask dilation. --- invokeai/app/invocations/vto_workflow.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/invokeai/app/invocations/vto_workflow.py b/invokeai/app/invocations/vto_workflow.py index 2589dbe286..e58c5473da 100644 --- a/invokeai/app/invocations/vto_workflow.py +++ b/invokeai/app/invocations/vto_workflow.py @@ -1,3 +1,4 @@ +import cv2 import numpy as np from PIL import Image @@ -32,6 +33,12 @@ class VTOInvocation(BaseInvocation, WithMetadata, WithBoard): ge=0.0, ) + mask_dilation: int = InputField( + description="The number of pixels to dilate the mask by. Default is 1.", + default=1, + ge=0, + ) + def invoke(self, context: InvocationContext) -> ImageOutput: # TODO(ryand): Avoid all the unnecessary flip-flopping between PIL and numpy. original_image = context.images.get_pil(self.original_image.image_name) @@ -59,9 +66,13 @@ class VTOInvocation(BaseInvocation, WithMetadata, WithBoard): pattern_with_shadows = multiply_images(expanded_pattern, Image.fromarray(shadows)) + # Dilate the mask. + clothing_mask_np = np.array(clothing_mask) + if self.mask_dilation > 0: + clothing_mask_np = cv2.dilate(clothing_mask_np, np.ones((3, 3), np.uint8), iterations=self.mask_dilation) + # Merge the pattern with the model image. pattern_with_shadows_np = np.array(pattern_with_shadows) - clothing_mask_np = np.array(clothing_mask) original_image_np = np.array(original_image) merged_image = np.where(clothing_mask_np[:, :, None], pattern_with_shadows_np, original_image_np) merged_image = Image.fromarray(merged_image)