watermark on rectangles

utils.py

@@ -17,7 +17,7 @@ apply_canny = CannyDetector()
 apply_openpose = OpenposeDetector()
 
 
-def add_watermark(image, im_size, watermark_path="__assets__/picsart_watermark.jpg",
+def add_watermark(image, im_size_h, im_size_w, watermark_path="__assets__/picsart_watermark.jpg",
                   wmsize=16, bbuf=5, opacity=0.9):
     '''
     Creates a watermark on the saved inference image.
@@ -25,8 +25,9 @@ def add_watermark(image, im_size, watermark_path="__assets__/picsart_watermark.j
     Shi-Lab's work.
     '''
     watermark = Image.open(watermark_path).resize((wmsize, wmsize))
-    loc = im_size - wmsize - bbuf
-    image[loc:-bbuf, loc:-bbuf, :] = watermark
+    loc_h = im_size_h - wmsize - bbuf
+    loc_w = im_size_w - wmsize - bbuf
+    image[loc_h:-bbuf, loc_w:-bbuf, :] = watermark
     return image
 
 
@@ -72,7 +73,9 @@ def create_video(frames, fps, rescale=False, path=None):
         if rescale:
             x = (x + 1.0) / 2.0  # -1,1 -> 0,1
         x = (x * 255).numpy().astype(np.uint8)
-        x = add_watermark(x, im_size=512)
+
+        h_, w_, _ = x.shape
+        x = add_watermark(x, im_size_h=h_, im_size_w=w_)
         outputs.append(x)
         # imageio.imsave(os.path.join(dir, os.path.splitext(name)[0] + f'_{i}.jpg'), x)
 
@@ -90,7 +93,8 @@ def create_gif(frames, fps, rescale=False):
         if rescale:
             x = (x + 1.0) / 2.0  # -1,1 -> 0,1
         x = (x * 255).numpy().astype(np.uint8)
-        x = add_watermark(x, im_size=512)
+        h_, w_, _ = x.shape
+        x = add_watermark(x, im_size_h=h_, im_size_w=w_)
         outputs.append(x)
         # imageio.imsave(os.path.join(dir, os.path.splitext(name)[0] + f'_{i}.jpg'), x)