feat: support meta SSL watermarking

SSL_watermark.py

@@ -0,0 +1,87 @@
+import numpy as np
+import json
+
+import torch
+from torchvision import transforms
+
+import torch_utils 
+import image_utils
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+torch.manual_seed(0)
+np.random.seed(0)
+
+print('Building backbone and normalization layer...')
+backbone = torch_utils.build_backbone(path='dino_r50.pth')
+normlayer = torch_utils.load_normalization_layer(path='out2048.pth')
+model = torch_utils.NormLayerWrapper(backbone, normlayer)
+
+print('Building the hypercone...')
+FPR = 1e-6
+angle = 1.462771101178447 # value for FPR=1e-6 and D=2048
+rho = 1 + np.tan(angle)**2
+carrier = torch.randn(1, 2048)
+carrier /= torch.norm(carrier, dim=1, keepdim=True)
+
+default_transform = transforms.Compose([
+        transforms.ToTensor(), 
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+
+def encode(image, epochs=10, psnr=44, lambda_w=1, lambda_i=1):
+    img_orig = default_transform(image).to(device, non_blocking=True).unsqueeze(0)
+    img = img_orig.clone().to(device, non_blocking=True) 
+    img.requires_grad = True
+    optimizer = torch.optim.Adam([img], lr=1e-2)
+
+    for iteration in range(epochs):
+        print(f'iteration: {iteration}')
+        x = image_utils.ssim_attenuation(img, img_orig)
+        x = image_utils.psnr_clip(x, img_orig, psnr)
+
+        ft = model(x) # BxCxWxH -> BxD
+
+        dot_product = (ft @ carrier.T) # BxD @ Dx1 -> Bx1
+        norm = torch.norm(ft, dim=-1, keepdim=True) # Bx1
+        cosines = torch.abs(dot_product/norm)
+        log10_pvalue = np.log10(torch_utils.cosine_pvalue(cosines.item(), ft.shape[-1]))
+        loss_R = -(rho * dot_product**2 - norm**2) # B-B -> B
+
+        loss_l2_img = torch.norm(x - img_orig)**2 # CxWxH -> 1
+        loss = lambda_w*loss_R + lambda_i*loss_l2_img
+        
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        logs = {
+            "keyword": "img_optim",
+            "iteration": iteration,
+            "loss": loss.item(),
+            "loss_R": loss_R.item(),
+            "loss_l2_img": loss_l2_img.item(),
+            "log10_pvalue": log10_pvalue.item(),
+        }
+        print("__log__:%s" % json.dumps(logs))
+
+    img = image_utils.ssim_attenuation(img, img_orig)
+    img = image_utils.psnr_clip(img, img_orig, psnr)
+    img = image_utils.round_pixel(img)
+    img = img.squeeze(0).detach().cpu()
+    img = transforms.ToPILImage()(image_utils.unnormalize_img(img).squeeze(0))
+
+    return img
+
+def decode(image):
+    img = default_transform(image).to(device, non_blocking=True).unsqueeze(0)
+    ft = model(img) # BxCxWxH -> BxD
+
+    dot_product = (ft @ carrier.T) # BxD @ Dx1 -> Bx1
+    norm = torch.norm(ft, dim=-1, keepdim=True) # Bx1
+    cosines = torch.abs(dot_product/norm)
+    log10_pvalue = np.log10(torch_utils.cosine_pvalue(cosines.item(), ft.shape[-1]))
+    loss_R = -(rho * dot_product**2 - norm**2) # B-B -> B
+
+    text_marked = "marked" if loss_R < 0 else "unmarked"
+    return f'Image is {text_marked}, with p-value={10**log10_pvalue}'

app.py

@@ -1,6 +1,7 @@
 import gradio as gr 
 from steganography import Steganography
 from utils import draw_multiple_line_text, generate_qr_code
+from SSL_watermark import encode, decode
 
 
 TITLE = """<h2 align="center"> ✍️ Invisible Watermark </h2>"""
@@ -8,20 +9,27 @@ TITLE = """<h2 align="center"> ✍️ Invisible Watermark </h2>"""
 
 def apply_watermark(radio_button, input_image, watermark_image, watermark_text, watermark_url):
     input_image = input_image.convert('RGB')
-
+    print(f'radio_button: {radio_button}')
     if radio_button == "Image":
         watermark_image = watermark_image.resize((input_image.width, input_image.height)).convert('L').convert('RGB')
         return Steganography().merge(input_image, watermark_image, digit=7)
     elif radio_button == "Text":
         watermark_image = draw_multiple_line_text(input_image.size, watermark_text)
         return Steganography().merge(input_image, watermark_image, digit=7)
-    else:
+    elif radio_button == "QRCode":
         size = min(input_image.width, input_image.height)
         watermark_image = generate_qr_code(watermark_url).resize((size, size)).convert('RGB')
         return Steganography().merge(input_image, watermark_image, digit=7)
+    else:
+        print('start encoding ssl watermark...')
+        return encode(input_image, epochs=5)
 
-def extract_watermark(input_image_to_extract):
-    return Steganography().unmerge(input_image_to_extract.convert('RGB'), digit=7).convert('RGBA')
+def extract_watermark(extract_radio_button, input_image_to_extract):
+    if extract_radio_button == 'Steganography':
+        return Steganography().unmerge(input_image_to_extract.convert('RGB'), digit=7).convert('RGBA')
+    else:
+        decoded_info = decode(image=input_image_to_extract)
+        return draw_multiple_line_text(input_image_size=input_image_to_extract.size, text=decoded_info)
 
 
 with gr.Blocks() as demo:
@@ -34,7 +42,7 @@ with gr.Blocks() as demo:
                 with gr.Blocks():
                     gr.Markdown("### Which type of watermark you want to apply?")
                     radio_button = gr.Radio(
-                        choices=["QRCode", "Text", "Image"], 
+                        choices=["QRCode", "Text", "Image", "SSL Watermark"], 
                         label="Watermark type", 
                         value="QRCode",
                         # info="Which type of watermark you want to apply?"
@@ -82,6 +90,11 @@ with gr.Blocks() as demo:
             with gr.Column():
                 gr.Markdown("### Image to extract watermark")
                 input_image_to_extract = gr.Image(type='pil')
+                extract_radio_button = gr.Radio(
+                    choices=["Steganography", "SSL Watermark"], 
+                    label="Extract methods", 
+                    value="Steganography"
+                )
             with gr.Column():
                 gr.Markdown("### Extracted watermark")
                 extracted_watermark = gr.Image(type='pil')
@@ -97,6 +110,10 @@ with gr.Blocks() as demo:
         inputs=[radio_button, input_image, watermark_image, watermark_text, watermark_url], 
         outputs=[output_image]
     )
-    extract_button.click(fn=extract_watermark, inputs=[input_image_to_extract], outputs=[extracted_watermark])
+    extract_button.click(
+        fn=extract_watermark, 
+        inputs=[extract_radio_button, input_image_to_extract], 
+        outputs=[extracted_watermark]
+    )
 
 demo.launch()

dino_r50.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ab26d85d00cb1be8e757cf8820cf0fd8aa729ea7e21b1cf6c44875952ba8eb0f
+size 788803344

image_utils.py

@@ -0,0 +1,80 @@
+import numpy as np
+
+import torch
+from torchvision import transforms
+
+import torch.nn.functional as F
+
+from torch.autograd.variable import Variable
+
+NORMALIZE_IMAGENET = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+image_mean = torch.Tensor(NORMALIZE_IMAGENET.mean).view(-1, 1, 1).to(device)
+image_std = torch.Tensor(NORMALIZE_IMAGENET.std).view(-1, 1, 1).to(device)
+
+def normalize_img(x):
+    return (x.to(device) - image_mean) / image_std
+
+def unnormalize_img(x):
+    return (x.to(device) * image_std) + image_mean
+
+def round_pixel(x):
+    x_pixel = 255 * unnormalize_img(x)
+    y = torch.round(x_pixel).clamp(0, 255)
+    y = normalize_img(y/255.0)
+    return y
+
+def project_linf(x, y, radius):
+    """ Clamp x-y so that Linf(x,y)<=radius """
+    delta = x - y
+    delta = 255 * (delta * image_std)
+    delta = torch.clamp(delta, -radius, radius)
+    delta = (delta / 255.0) / image_std
+    return y + delta
+
+def psnr_clip(x, y, target_psnr):
+    """ Clip x-y so that PSNR(x,y)=target_psnr """
+    delta = x - y
+    delta = 255 * (delta * image_std)
+    psnr = 20*np.log10(255) - 10*torch.log10(torch.mean(delta**2))
+    if psnr<target_psnr:
+        delta = (torch.sqrt(10**((psnr-target_psnr)/10))) * delta 
+    psnr = 20*np.log10(255) - 10*torch.log10(torch.mean(delta**2))
+    delta = (delta / 255.0) / image_std
+    return y + delta
+
+def ssim_heatmap(img1, img2, window_size):
+    """ Compute the SSIM heatmap between 2 images """
+    _1D_window = torch.Tensor(
+        [np.exp(-(x - window_size//2)**2/float(2*1.5**2)) for x in range(window_size)]
+        ).to(device, non_blocking=True)
+    _1D_window = (_1D_window/_1D_window.sum()).unsqueeze(1)
+    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+    window = Variable(_2D_window.expand(3, 1, window_size, window_size).contiguous())
+
+    mu1 = F.conv2d(img1, window, padding = window_size//2, groups = 3)
+    mu2 = F.conv2d(img2, window, padding = window_size//2, groups = 3)
+
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1*mu2
+
+    sigma1_sq = F.conv2d(img1*img1, window, padding = window_size//2, groups = 3) - mu1_sq
+    sigma2_sq = F.conv2d(img2*img2, window, padding = window_size//2, groups = 3) - mu2_sq
+    sigma12 = F.conv2d(img1*img2, window, padding = window_size//2, groups = 3) - mu1_mu2
+
+    C1 = 0.01**2
+    C2 = 0.03**2
+
+    ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))
+    return ssim_map
+
+def ssim_attenuation(x, y):
+    """ attenuate x-y using SSIM heatmap """
+    delta = x - y
+    ssim_map = ssim_heatmap(x, y, window_size=17) # 1xCxHxW
+    ssim_map = torch.sum(ssim_map, dim=1, keepdim=True)
+    ssim_map = torch.clamp_min(ssim_map,0)
+    delta = delta*ssim_map
+    return y + delta

out2048.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4b256188454d8f7cf440de048df398e2a3209136a52cd7cdac834f5792f526a3
+size 16786561

requirements.txt

@@ -1,4 +1,8 @@
+torch==1.10.1
+torchvision==0.11.2
 Pillow
 click
 gradio
 qrcode
+scipy
+json

torch_utils.py

@@ -0,0 +1,84 @@
+import numpy as np
+
+import torch
+import torch.nn as nn
+from torchvision import models
+
+from scipy.optimize import root_scalar
+from scipy.special import betainc
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def build_backbone(path, name='resnet50'):
+    """ Builds a pretrained ResNet-50 backbone. """
+    model = getattr(models, name)(pretrained=False)
+    model.head = nn.Identity()
+    model.fc = nn.Identity()
+    checkpoint = torch.load(path, map_location=device)
+    state_dict = checkpoint
+    for ckpt_key in ['state_dict', 'model_state_dict', 'teacher']:
+        if ckpt_key in checkpoint:
+            state_dict = checkpoint[ckpt_key]
+    state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+    state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()}
+    msg = model.load_state_dict(state_dict, strict=False)
+    return model
+
+def get_linear_layer(weight, bias):
+    """ Creates a layer that performs feature whitening or centering """
+    dim_out, dim_in = weight.shape
+    layer = nn.Linear(dim_in, dim_out)
+    layer.weight = nn.Parameter(weight)
+    layer.bias = nn.Parameter(bias)
+    return layer
+
+def load_normalization_layer(path):
+    """
+    Loads the normalization layer from a checkpoint and returns the layer.
+    """
+    checkpoint = torch.load(path, map_location=device)
+    if 'whitening' in path or 'out' in path:
+        D = checkpoint['weight'].shape[1]
+        weight = torch.nn.Parameter(D*checkpoint['weight'])
+        bias = torch.nn.Parameter(D*checkpoint['bias'])
+    else:
+        weight = checkpoint['weight']
+        bias = checkpoint['bias']
+    return get_linear_layer(weight, bias).to(device, non_blocking=True)
+
+class NormLayerWrapper(nn.Module):
+    """
+    Wraps backbone model and normalization layer
+    """
+    def __init__(self, backbone, head):
+        super(NormLayerWrapper, self).__init__()
+        backbone.eval(), head.eval()
+        self.backbone = backbone
+        self.head = head
+
+    def forward(self, x):
+        output = self.backbone(x)
+        return self.head(output)
+
+def cosine_pvalue(c, d, k=1):
+    """
+    Returns the probability that the absolute value of the projection
+    between random unit vectors is higher than c
+    Args:
+        c: cosine value
+        d: dimension of the features
+        k: number of dimensions of the projection
+    """
+    assert k>0
+    a = (d - k) / 2.0
+    b = k / 2.0
+    if c < 0:
+        return 1.0
+    return betainc(a, b, 1 - c ** 2)
+
+def pvalue_angle(dim, k=1, angle=None, proba=None):
+    def f(a):
+        return cosine_pvalue(np.cos(a), dim, k) - proba
+    a = root_scalar(f, x0=0.49*np.pi, bracket=[0, np.pi/2])
+    # a = fsolve(f, x0=0.49*np.pi)[0]
+    return a.root