watermarklab 0.1.1

A comprehensive toolkit for robust image watermarking benchmarking and development.

WatermarkLab

WatermarkLab is a powerful toolkit for robust image watermarking research and development. It provides a complete suite of tools for watermark embedding, extraction, robustness testing, and performance evaluation, helping researchers and developers easily implement and evaluate robust image watermarking algorithms.

Installation

Install WatermarkLab via pip:

pip install watermarklab

Download Resources

huggingface-cli download chenoly/watermarklab
huggingface-cli download stabilityai/stable-diffusion-2-1-base
huggingface-cli download stabilityai/stable-diffusion-2-1

Benchmarking PGWs

import watermarklab as wl
from watermarklab.watermarks.PGWs import VINE
from watermarklab.utils.data import DataLoader
from watermarklab.datasets import MS_COCO_2017_VAL_IMAGES
from watermarklab.attackers.attackerloader import AttackersWithFactorsModel
from watermarklab.watermarks.PGWs import rivaGAN, dctDwtSvd, TrustMark, StegaStamp, InvisMark
import warnings

warnings.filterwarnings('ignore')


default_attackers = AttackersWithFactorsModel()

mscoco2017_256_32 = MS_COCO_2017_VAL_IMAGES(im_size=256, bit_len=32)
dataloader_256_32 = DataLoader(mscoco2017_256_32, batch_size=32)

rivagan = rivaGAN(bits_len=32, img_size=256)
wl.evaluate("save_results/PGWs", rivagan, default_attackers, dataloader_256_32, mode='decode', noise_save=True)

dctdwtsvd = dctDwtSvd(bits_len=32, img_size=256)
wl.evaluate("save_results/PGWs", dctdwtsvd, default_attackers, dataloader_256_32, mode='decode', noise_save=True)

mscoco2017_256_100 = MS_COCO_2017_VAL_IMAGES(im_size=256, bit_len=100)
dataloader_256_100 = DataLoader(mscoco2017_256_100, batch_size=32)
invismark = InvisMark(bits_len=100, img_size=256)
wl.evaluate("save_results/PGWs", invismark, default_attackers, dataloader_256_100, mode='decode', noise_save=True)

trustmark = TrustMark(bits_len=100, img_size=256)
wl.evaluate("save_results/PGWs", trustmark, default_attackers, dataloader_256_100, mode='decode', noise_save=True)

mscoco2017_400_100 = MS_COCO_2017_VAL_IMAGES(im_size=400, bit_len=100)
dataloader_400_100 = DataLoader(mscoco2017_400_100, batch_size=16)
stegastamp = StegaStamp(bits_len=100, img_size=400)
wl.evaluate("save_results/PGWs", stegastamp, default_attackers, dataloader_400_100, mode='decode', noise_save=True)


mscoco2017_256_100 = MS_COCO_2017_VAL_IMAGES(im_size=256, bit_len=100)
dataloader_256_100 = DataLoader(mscoco2017_256_100, batch_size=32)
vine = VINE(bits_len=100, img_size=256)
wl.evaluate("save_results/PGWs", vine, default_attackers, dataloader_256_100, noise_save=True)

Benchmarking IGWs

import watermarklab as wl
from watermarklab.utils.data import DataLoader
from watermarklab.datasets import MS_COCO_2017_VAL_PROMPTS
from watermarklab.attackers.attackerloader import AttackersWithFactorsModel
from watermarklab.watermarks.IGWs import StableSignature, GaussianShading, TreeRing
import warnings


warnings.filterwarnings('ignore')
default_attackers = AttackersWithFactorsModel()

gaussianshading = GaussianShading(local_files_only=True)
mscoco_256 = MS_COCO_2017_VAL_PROMPTS(bit_len=256)
wl.evaluate("save_results/IGWs/", gaussianshading, default_attackers, DataLoader(mscoco_256, batch_size=128), noise_save=True)

treering = TreeRing(local_files_only=True)
mscoco_1 = MS_COCO_2017_VAL_PROMPTS(bit_len=1)
wl.evaluate("save_results/IGWs/", treering, default_attackers, DataLoader(mscoco_1, batch_size=32), noise_save=True)

stablesignature = StableSignature(local_files_only=True)
mscoco_48 = MS_COCO_2017_VAL_PROMPTS(bit_len=48)
wl.evaluate("save_results/IGWs/", stablesignature, default_attackers, DataLoader(mscoco_48, batch_size=32), need_cover=True, noise_save=True)```

How to visualize results?

import watermarklab as wl
from watermarklab.attackers.attackerloader import AttackersWithFactorsModel

results = ['saved_all_json/result_dctDwtSvd.json',
           'saved_all_json/result_rivaGAN.json',
           'saved_all_json/result_TrustMark-Q.json',
           'saved_all_json/result_InvisMark.json',
           'saved_all_json/result_StegaStamp.json',
           'saved_all_json/result_VINE-R.json',
           'saved_all_json/result_GaussianShading.json',
           'saved_all_json/result_StableSignature.json',
           'saved_all_json/result_TreeRing-Ring.json']

default_attackers = AttackersWithFactorsModel()
wl.draw.plot_model_robustness_under_single_attack(results, "draw_ALL/MR_SA")
wl.draw.plot_model_robustness_under_all_attack(results, "draw_ALL/MR_AA")
wl.draw.plot_model_robustness_scores_under_all_attacks(results, "draw_ALL/MRS_AA")
wl.draw.plot_model_robustness_ranking_under_single_attack(results, "draw_ALL/MRK_SA")
wl.draw.test_compute_overall_robustness_scores(results)
wl.draw.plot_model_robustness_ranking_by_attacker_group(results, "draw_ALL/MRK_GA", default_attackers.attacker_groups)
wl.draw.plot_model_overall_robustness_ranking(results, "draw_ALL/MRK_MO")
wl.draw.plot_all_attack_ranking(results, "draw_ALL/ARK")
wl.draw.plot_attack_effectiveness_at_tpr_levels(results, "draw_ALL/ARK_TPR", tpr_levels=[0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1])
wl.draw.plot_visual_quality(results, "draw_ALL/VQ", show_dataset_name=False)
wl.draw.plot_stego_visualization(results, "draw_ALL/MVC")
wl.draw.plot_attack_visualization(results, "draw_ALL/MVC")

License

WatermarkLab is licensed under the MIT License. See the license file for details.