Spaces:
Runtime error
Runtime error
File size: 2,006 Bytes
f2dbf59 |
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 |
import torch
import numpy as np
from PIL import Image, ImageSequence, ImageOps
from ..utils import tensor2pil, pil2tensor
PADDING = 4
class WatermarkNode:
@classmethod
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
|