fix: 🐛 batch support

This commit is contained in:
melMass
2023-10-09 17:40:52 +02:00
parent c1d42de0fc
commit f1ff9fc7c4
2 changed files with 26 additions and 24 deletions
+20 -21
View File
@@ -13,6 +13,7 @@ from PIL.PngImagePlugin import PngInfo
from skimage.filters import gaussian
from skimage.util import compare_images
from ..log import log
from ..utils import pil2tensor, tensor2np, tensor2pil
# try:
@@ -318,22 +319,26 @@ class MaskToImage:
FUNCTION = "render_mask"
def render_mask(self, mask, color, background):
mask = tensor2np(mask)
mask = Image.fromarray(mask).convert("L")
masks = tensor2np(mask)
images = []
for m in masks:
_mask = Image.fromarray(m).convert("L")
image = Image.new("RGBA", mask.size, color=color)
# apply the mask
image = Image.composite(
image, Image.new("RGBA", mask.size, color=background), mask
)
log.debug(f"Converted mask to PIL Image format, size: {_mask.size}")
# image = ImageChops.multiply(image, mask)
# apply over background
# image = Image.alpha_composite(Image.new("RGBA", image.size, color=background), image)
image = Image.new("RGBA", _mask.size, color=color)
# apply the mask
image = Image.composite(
image, Image.new("RGBA", _mask.size, color=background), _mask
)
image = pil2tensor(image.convert("RGB"))
# image = ImageChops.multiply(image, mask)
# apply over background
# image = Image.alpha_composite(Image.new("RGBA", image.size, color=background), image)
return (image,)
images.append(image.convert("RGB"))
return (pil2tensor(images),)
class ColoredImage:
@@ -381,19 +386,13 @@ class ImagePremultiply:
CATEGORY = "mtb/image"
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("RGBA",)
FUNCTION = "premultiply"
def premultiply(self, image, mask, invert):
images = tensor2pil(image)
if invert:
masks = tensor2pil(mask) # .convert("L")
else:
masks = tensor2pil(1.0 - mask)
single = False
if len(mask) == 1:
single = True
masks = tensor2pil(mask) if invert else tensor2pil(1.0 - mask)
single = len(mask) == 1
masks = [x.convert("L") for x in masks]
out = []
+6 -3
View File
@@ -1,7 +1,8 @@
from rembg import remove
from ..utils import pil2tensor, tensor2pil
from PIL import Image
import comfy.utils
from PIL import Image
from rembg import remove
from ..utils import pil2tensor, tensor2pil
class ImageRemoveBackgroundRembg:
@@ -91,6 +92,8 @@ class ImageRemoveBackgroundRembg:
image_on_bg.paste(img_rm, mask=mask)
image_on_bg = image_on_bg.convert("RGB")
out_img.append(img_rm)
out_mask.append(mask)
out_img_on_bg.append(image_on_bg)