Spaces:
Runtime error
Runtime error
| import torch | |
| import numpy as np | |
| from PIL import Image, ImageSequence, ImageOps | |
| from ..utils import tensor2pil, pil2tensor | |
| PADDING = 4 | |
| class WatermarkNode: | |
| def INPUT_TYPES(cls): | |
| return { | |
| "required": { | |
| "images": ("IMAGE",), | |
| "logo_list": ("IMAGE",), | |
| }, | |
| "optional": { | |
| "logo_mask": ("MASK",), | |
| "enabled": ("BOOLEAN", {"default": True}),} | |
| } | |
| RETURN_TYPES = ("IMAGE",) | |
| FUNCTION = "watermark" | |
| CATEGORY = "tbox/Image" | |
| def watermark(self, images, logo_list, logo_mask, enabled): | |
| outputs = [] | |
| if enabled == False: | |
| return(images,) | |
| print(f'logo shape: {logo_list.shape}') | |
| print(f'images shape: {images.shape}') | |
| logo = tensor2pil(logo_list[0]) | |
| if logo_mask is not None: | |
| logo_mask = tensor2pil(logo_mask) | |
| for i, image in enumerate(images): | |
| img = tensor2pil(image) #Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) | |
| dst = self.add_watermark2(img, logo, logo_mask, 85) | |
| result = pil2tensor(dst) | |
| outputs.append(result) | |
| base_image = torch.stack([tensor.squeeze() for tensor in outputs]) | |
| return (base_image,) | |
| def add_watermark2(self, image, logo, logo_mask, opacity=None): | |
| logo_width, logo_height = logo.size | |
| image_width, image_height = image.size | |
| if image_height <= logo_height + PADDING * 2 or image_width <= logo_width + PADDING * 2: | |
| return image | |
| y = image_height - logo_height - PADDING * 1 | |
| x = PADDING | |
| logo = logo.convert('RGBA') | |
| opacity = int(opacity / 100 * 255) | |
| logo.putalpha(Image.new("L", logo.size, opacity)) | |
| if logo_mask is not None: | |
| logo.putalpha(ImageOps.invert(logo_mask)) | |
| position = (x, y) | |
| image.paste(logo, position, logo) | |
| return image | |