watermark

utils.py

@@ -6,6 +6,7 @@ from torchvision.transforms import Resize
 import imageio
 from einops import rearrange
 import cv2
+from PIL import Image
 from annotator.util import resize_image, HWC3
 from annotator.canny import CannyDetector
 from annotator.openpose import OpenposeDetector
@@ -16,7 +17,7 @@ apply_canny = CannyDetector()
 apply_openpose = OpenposeDetector()
 
 
-def add_watermark(image, im_size, watermark_path="__assets__/pair_watermark.png",
+def add_watermark(image, im_size, watermark_path="__assets__/picsart_watermark.jpg",
                   wmsize=16, bbuf=5, opacity=0.9):
     '''
     Creates a watermark on the saved inference image.
@@ -25,7 +26,7 @@ def add_watermark(image, im_size, watermark_path="__assets__/pair_watermark.png"
     '''
     watermark = Image.open(watermark_path).resize((wmsize, wmsize))
     loc = im_size - wmsize - bbuf
-    image[:,:,loc:-bbuf, loc:-bbuf] = watermark
+    image[loc:-bbuf, loc:-bbuf, :] = watermark
     return image