mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
feat(nodes): add alpha mask to tensor invocation
This commit is contained in:
parent
c89a24d1ea
commit
fc26f3e430
@ -1,7 +1,8 @@
|
|||||||
|
import numpy as np
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from invokeai.app.invocations.baseinvocation import BaseInvocation, InvocationContext, invocation
|
from invokeai.app.invocations.baseinvocation import BaseInvocation, InvocationContext, invocation
|
||||||
from invokeai.app.invocations.fields import InputField, TensorField, WithMetadata
|
from invokeai.app.invocations.fields import ImageField, InputField, TensorField, WithMetadata
|
||||||
from invokeai.app.invocations.primitives import MaskOutput
|
from invokeai.app.invocations.primitives import MaskOutput
|
||||||
|
|
||||||
|
|
||||||
@ -34,3 +35,25 @@ class RectangleMaskInvocation(BaseInvocation, WithMetadata):
|
|||||||
width=self.width,
|
width=self.width,
|
||||||
height=self.height,
|
height=self.height,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@invocation(
|
||||||
|
"alpha_mask_to_tensor",
|
||||||
|
title="Alpha Mask to Tensor",
|
||||||
|
tags=["conditioning"],
|
||||||
|
category="conditioning",
|
||||||
|
version="1.0.0",
|
||||||
|
)
|
||||||
|
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.")
|
||||||
|
|
||||||
|
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)
|
||||||
|
|
||||||
|
return MaskOutput(
|
||||||
|
mask=TensorField(tensor_name=context.tensors.save(mask)), height=image.height, width=image.width
|
||||||
|
)
|
||||||
|
Loading…
Reference in New Issue
Block a user