Watermark encoder expects images in BGR channel order (matching cv2 i…

…mread). This fix reduces the watermark artifacts.

scripts/demo/detect.py

@@ -83,7 +83,7 @@ def __init__(self, watermark):
     def __call__(self, x: np.ndarray) -> np.ndarray:
         """
         Detects the number of matching bits the predefined watermark with one
-        or multiple images. Images should be in cv2 format, e.g. h x w x c.
+        or multiple images. Images should be in cv2 format, e.g. h x w x c BGR.
 
         Args:
             x: ([B], h w, c) in range [0, 255]
@@ -94,7 +94,6 @@ def __call__(self, x: np.ndarray) -> np.ndarray:
         squeeze = len(x.shape) == 3
         if squeeze:
             x = x[None, ...]
-        x = np.flip(x, axis=-1)
 
         bs = x.shape[0]
         detected = np.empty((bs, self.num_bits), dtype=bool)

scripts/demo/streamlit_helpers.py

@@ -43,19 +43,19 @@ def __call__(self, image: torch.Tensor):
         Returns:
             same as input but watermarked
         """
-        # watermarking libary expects input as cv2 format
+        # watermarking libary expects input as cv2 BGR format
         squeeze = len(image.shape) == 4
         if squeeze:
             image = image[None, ...]
         n = image.shape[0]
         image_np = rearrange(
             (255 * image).detach().cpu(), "n b c h w -> (n b) h w c"
-        ).numpy()
+        ).numpy()[:,:,:,::-1]
         # torch (b, c, h, w) in [0, 1] -> numpy (b, h, w, c) [0, 255]
         for k in range(image_np.shape[0]):
             image_np[k] = self.encoder.encode(image_np[k], "dwtDct")
         image = torch.from_numpy(
-            rearrange(image_np, "(n b) h w c -> n b c h w", n=n)
+            rearrange(image_np[:,:,:,::-1], "(n b) h w c -> n b c h w", n=n)
         ).to(image.device)
         image = torch.clamp(image / 255, min=0.0, max=1.0)
         if squeeze: