added encoding and decoding

app.py

@@ -3,80 +3,94 @@ import gradio.inputs as grinputs
 import gradio.outputs as groutputs
 
 import numpy as np
+import json
 
 import torch
-import torch.nn as nn
-from torchvision import models
+from torchvision import transforms
+
+import utils 
+import utils_img
 
 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 = utils.build_backbone(path='dino_r50.pth')
+normlayer = utils.load_normalization_layer(path='out2048.pth')
+model = utils.NormLayerWrapper(backbone, normlayer)
+
+print('Building the hypercone...')
 FPR = 1e-6
-carrier = np.random.randn(1, 2048)
-
-
-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)
-
-backbone = build_backbone(path='dino_r50.pth')
-normlayer = load_normalization_layer(path='out2048.pth')
-model = NormLayerWrapper(backbone, normlayer)
-
-def encode(image):
-    return image
+angle = 1.462771101178447 # value for FPR=1e-6 and D=2048
+rho = 1 + np.tan(angle)**2
+# angle = utils.pvalue_angle(2048, 1, proba=FPR)
+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):
+        x = utils_img.ssim_attenuation(img, img_orig)
+        x = utils_img.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(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 = utils_img.ssim_attenuation(img, img_orig)
+    img = utils_img.psnr_clip(img, img_orig, psnr)
+    img = utils_img.round_pixel(img)
+    img = img.squeeze(0).detach().cpu()
+    img = transforms.ToPILImage()(utils_img.unnormalize_img(img).squeeze(0))
+
+    return img
 
 def decode(image):
-    return 'decoded'
+    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(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 'Image is {s}, with p-value={p}'.format(s=text_marked, p=10**log10_pvalue)
+
+
 
 def on_submit(image, mode):
     print('{} mode'.format(mode))

requirements.txt

@@ -1,3 +1,4 @@
 torch==1.10.1
 torchvision==0.11.2
 pillow==9.0.0
+scipy

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

utils_img.py

@@ -0,0 +1,85 @@
+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)
+    # min_v = torch.min(ssim_map)
+    # range_v = torch.max(ssim_map) - min_v
+    # if range_v < 1e-10:
+    #     return y + delta
+    # ssim_map = (ssim_map - min_v) / range_v
+    delta = delta*ssim_map
+    return y + delta