mask_operations.py

from PIL import ImageOps, ImageChops
from typing import Literal

from invokeai.invocation_api import (
    BaseInvocation,
    InputField,
    invocation,
    InvocationContext,
    ImageField,
    ImageOutput
)

OP_MODES = Literal[
    "OR",
    "SUB",
    "AND"
]

@invocation(
    "mask_operations",
    title="Mask Operations",
    tags=["image", "mask"],
    category="image",
    version="1.0.0",
)
class MaskOperationsIvocation(BaseInvocation):
    """Mask Operations"""
    mask_a: ImageField = InputField(default=None, description="A mask")
    mask_b: ImageField = InputField(default=None, description="B mask")
    operation: OP_MODES = InputField(default="OR", description="Operation")
    invert_a: bool = InputField(default=False, description="Invert A input mask before the operation")
    invert_b: bool = InputField(default=False, description="Invert b input mask before the operation")
    invert_output: bool = InputField(default=False, description="Inverted output")

    def invoke(self, context: InvocationContext) -> ImageOutput:
        mask_a = context.images.get_pil(self.mask_a.image_name).convert('L')
        mask_b = context.images.get_pil(self.mask_b.image_name).convert('L')
 
        if mask_a.size != mask_b.size:
            mask_b = mask_b.resize(mask_a.size)
        
        if self.invert_a:
            mask_a = ImageOps.invert(mask_a)

        if self.invert_b:
            mask_b = ImageOps.invert(mask_b)


        if self.operation == "OR":
            result = ImageChops.lighter(mask_a, mask_b)

        if self.operation == "SUB":
            result = ImageChops.subtract(mask_a, mask_b, scale=1.0, offset=0)

        if self.operation == "AND":
            result = ImageChops.darker(mask_a, mask_b)


        if self.invert_output:
            result = ImageOps.invert(result)

        image_dto = context.images.save(image=result.convert('RGBA'))

        return ImageOutput.build(image_dto)