Make MaskToImage support masks with a batch size.
This commit is contained in:
parent
8e3d1cbf3b
commit
51583164ef
|
@ -72,7 +72,7 @@ class MaskToImage:
|
||||||
FUNCTION = "mask_to_image"
|
FUNCTION = "mask_to_image"
|
||||||
|
|
||||||
def mask_to_image(self, mask):
|
def mask_to_image(self, mask):
|
||||||
result = mask[None, :, :, None].expand(-1, -1, -1, 3)
|
result = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
|
||||||
return (result,)
|
return (result,)
|
||||||
|
|
||||||
class ImageToMask:
|
class ImageToMask:
|
||||||
|
|
Loading…
Reference in New Issue