About baseline method FNNS

[fenghe12]

Could you please share the codes for this baseline? In your paper, you mentioned FNNS as one of the baselines, but it was used to replace the original SteganoGan with HiDDeN. I found the implementation of FNNS online, but it seems that the embedded watermark is an image (e.g., a four-dimensional tensor of b c h w).

[pierrefdz]

Hello, here is what I used. It's not exactly similar to the original paper, I made some improvements. Notably steganogan does not handle crops well so I used HiDDeN and I optimize with Adam.

class FNNS(Watermark):

    def __init__(self, nbits, nsteps=10, data_aug="none", decoder_path="/path/to/wm/decoder", **kwargs):
        self.model = torch.jit.load(decoder_path).to(device)
        print("Model loaded")
        self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        self.unnormalize = transforms.Normalize(mean=[-0.485/0.229, -0.456/0.224, -0.406/0.225], std=[1/0.229, 1/0.224, 1/0.225])
        self.attenuation = JND(preprocess = self.unnormalize).to(device)
        self.data_aug = nn.Identity().to(device)
        self.nsteps = nsteps 
        if data_aug != "none":
            self.nsteps *= 10

    def loss_w(self, ft, msg, T=1e1):
        return F.binary_cross_entropy_with_logits(ft/T, msg)

    def encode(self, img:Image, msg:np.ndarray, scaling:float=3.0) -> Image:
        msg = torch.tensor(msg, dtype=torch.float, device=device) # k
        img = self.normalize(transforms.ToTensor()(img).unsqueeze(0).to(device)) # 1 3 h w
        heatmap = scaling * self.attenuation.heatmaps(img) # 1 1 h w
        # distortion_scaler = torch.tensor([0.072*(1/0.299), 0.072*(1/0.587), 0.072*(1/0.114)])[:,None,None].to(device)
        distortion_scaler = 1.0
        distortion = 1e-3 * torch.randn_like(img, device=device) # 1 c h w
        distortion.requires_grad = True
        optimizer = torch.optim.Adam([distortion], lr=0.5e-0)
        log_stats = []
        iter_time = time.time()
        for gd_it in range(self.nsteps):
            gd_it_time = time.time()
            clip_img = img + heatmap * torch.tanh(distortion) * distortion_scaler
            ft = self.model(self.data_aug(clip_img)) # 1 d
            loss_w = self.loss_w(ft[0], msg)
            optimizer.zero_grad()
            loss_w.backward()
            optimizer.step()
            with torch.no_grad():
                psnrs = utils_img.psnr(clip_img, img, img_space='img')
                log_stats.append({
                    'gd_it': gd_it,
                    'loss_w': loss_w.item(),
                    'psnr': torch.nanmean(psnrs).item(),
                    'gd_it_time': time.time() - gd_it_time,
                    'iter_time': time.time() - iter_time,
                    'max_mem': torch.cuda.max_memory_allocated() / (1024*1024),
                    'kw': 'optim',
                })
                if (gd_it+1) % 5 == 0:
                    print(json.dumps(log_stats[-1]))
        clip_img = img + heatmap * torch.tanh(distortion) * distortion_scaler
        clip_img = torch.clamp(self.unnormalize(clip_img), 0, 1)
        clip_img = torch.round(255 * clip_img)/255 
        return transforms.ToPILImage()(clip_img.squeeze(0).cpu())

    @torch.no_grad()
    def decode(self, img: Image) -> np.ndarray:
        img = self.normalize(transforms.ToTensor()(img).unsqueeze(0).to(device)) # 1 3 h w
        ft = self.model(img) # 1 d
        return (ft > 0).squeeze().cpu().numpy()