fix(nodes): handle invert in alpha_mask_to_tensor

This commit is contained in:
psychedelicious 2024-04-15 12:37:30 +10:00 committed by Kent Keirsey
parent d1db6198b5
commit 602a59066e

View File

@ -48,12 +48,15 @@ class AlphaMaskToTensorInvocation(BaseInvocation):
"""Convert a mask image to a tensor. Opaque regions are 1 and transparent regions are 0."""
image: ImageField = InputField(description="The mask image to convert.")
invert: bool = InputField(default=False, description="Invert the mask (1s become 0s and 0s become 1s).")
invert: bool = InputField(default=False, description="Whether to invert the mask.")
def invoke(self, context: InvocationContext) -> MaskOutput:
image = context.images.get_pil(self.image.image_name)
mask = torch.zeros((1, image.height, image.width), dtype=torch.bool)
mask[0] = torch.tensor(np.array(image)[:, :, 3] > 0, dtype=torch.bool)
if self.invert:
mask[0] = torch.tensor(np.array(image)[:, :, 3] == 0, dtype=torch.bool)
else:
mask[0] = torch.tensor(np.array(image)[:, :, 3] > 0, dtype=torch.bool)
return MaskOutput(
mask=TensorField(tensor_name=context.tensors.save(mask)), height=image.height, width=image.width