File size: 1,929 Bytes
e00eceb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | import torch
from .imagefunc import log, pil2tensor,image2mask, extract_numbers
from PIL import Image
class BatchSelector:
def __init__(self):
self.NODE_NAME = 'BatchSelector'
pass
@classmethod
def INPUT_TYPES(self):
return {
"required": {
"select": ("STRING", {"default": "0,"},),
},
"optional": {
"images": ("IMAGE",), #
"masks": ("MASK",), #
}
}
RETURN_TYPES = ("IMAGE", "MASK",)
RETURN_NAMES = ("image", "mask",)
FUNCTION = 'batch_selector'
CATEGORY = '😺dzNodes/LayerUtility/SystemIO'
def batch_selector(self, select, images=None, masks=None
):
ret_images = []
ret_masks = []
empty_image = pil2tensor(Image.new("RGBA", (64, 64), (0, 0, 0, 0)))
empty_mask = image2mask(Image.new("L", (64, 64), color="black"))
indexs = extract_numbers(select)
for i in indexs:
if images is not None:
if i < len(images):
ret_images.append(images[i].unsqueeze(0))
else:
ret_images.append(images[-1].unsqueeze(0))
if masks is not None:
if i < len(masks):
ret_masks.append(masks[i].unsqueeze(0))
else:
ret_masks.append(masks[-1].unsqueeze(0))
if len(ret_images) == 0:
ret_images.append(empty_image)
if len(ret_masks) == 0:
ret_masks.append(empty_mask)
log(f"{self.NODE_NAME} Processed {len(ret_images)} image(s).", message_type='finish')
return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),)
NODE_CLASS_MAPPINGS = {
"LayerUtility: BatchSelector": BatchSelector
}
NODE_DISPLAY_NAME_MAPPINGS = {
"LayerUtility: BatchSelector": "LayerUtility: Batch Selector"
} |