Algorithms
How to implement a new watermarking algorithm
This guide explains how to implement a new watermarking algorithm wrapper for integration with the WIBE framework.
The wrapper system provides a standardized interface for various watermarking techniques.
For more examples, refer to the wibench.algorithms module.
Create your_wrapper.py file in user_plugins directory.
Implement the wrapper class ctor
from wibench.algorithms import BaseAlgorithmWrapper
class MyAlgorithmWrapper(BaseAlgorithmWrapper):
"""Wrapper for My Watermarking Algorithm"""
# Unique identifier for your algorithm (lowercase, no spaces).
# Not strictly required; by default, it is the same as the class name.
name = "my_algorithm"
def __init__(self, params: dict):
"""
Initialize the wrapper.
Args:
params: Dictionary of configuration parameters. You may cast params
to a dataclass or just use dictionaries.
"""
super().__init__(params)
# You may need to import some external code here.
# If this code is implemented in a package, for example:
from trustmark import TrustMark
# If the external code resides in an unorganized project, the simplest way
# to import it is via sys.path (example from hidden):
import sys
sys.path.append(str(Path(params["module_path"])))
from utils import (
load_options,
load_last_checkpoint
)
# etc.
Implement the “watermark_data_gen” function
The watermark_data_gen function should provide any additional data the watermarking algorithm may require — for example, a bit message or watermark key. If the algorithm only requires an object to embed the watermark and a bit message sequence, you can use TorchBitWatermarkData:
from wibench.watermark_data import TorchBitWatermarkData
class MyAlgorithmWrapper(BaseAlgorithmWrapper):
...
def watermark_data_gen(self) -> TorchBitWatermarkData:
return TorchBitWatermarkData.get_random(self.params.wm_length) # or another parameter defining watermark length
Implement the “embed” function
This function embeds a watermark into the input object and returns the object with the watermark. For example, for image-based algorithms:
from wibench.watermark_data import TorchBitWatermarkData
from wibench.typing import TorchImg
class MyAlgorithmWrapper(BaseAlgorithmWrapper):
...
def embed(self, image: TorchImg, watermark_data: TorchBitWatermarkData) -> TorchImg:
...
You may find the following functions from wibench.utils useful:
- wibench.utils.torch_img2numpy_bgr(image: TorchImg) ndarray[source]
Convert torch image tensor to numpy array in BGR format.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format with values in [0, 1] range, float32
Returns
- np.ndarray
Output array in (H, W, C) BGR format with values in [0, 255] range, uint8
- wibench.utils.numpy_bgr2torch_img(image: ndarray) TorchImg[source]
Convert numpy BGR array to torch image tensor.
Parameters
- imagenp.ndarray
Input array in (H, W, C) BGR format with values in [0, 255] range, uint8
Returns
- TorchImg
Output tensor in (C, H, W) RGB format with values in [0, 1] range, float
- wibench.utils.resize_torch_img(image: TorchImg, size: List[int], mode: str = 'bilinear', align_corners: bool = True) TorchImg[source]
Resize a torch image tensor to specified dimensions.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- sizeList[int]
Target size as [height, width]
- modestr, optional
Interpolation mode (‘nearest’, ‘bilinear’, ‘bicubic’) Default is ‘bilinear’
- align_cornersbool, optional
Flag for align_corners parameter in interpolation Default is True
Returns
- TorchImg
Resized image tensor
- wibench.utils.overlay_difference(original_image: TorchImg, resized_image: TorchImg, marked_image: TorchImg, factor: float = 1.0) TorchImg[source]
Overlay difference between images of one size to image of another size.
Parameters
- original_imageTorchImg
Base reference image to overlay on
- resized_imageTorchImg
Resized version of original (should match marked_image size)
- marked_imageTorchImg
Watermarked or processed image
- factorfloat
Factor to enhance difference
Returns
- TorchImg
Original image with overlay
Notes
Computes difference between marked and resized images
Resizes difference to match original image dimensions
Adds difference to original image
- wibench.utils.normalize_image(image: TorchImg, transform: Optional[Normalize] = None) TorchImgNormalize[source]
Normalize image tensor from [0,1] to [-1,1] range and add batch dimension.
Parameters
- imageTorchImg
Input image in [0,1] range (C, H, W)
- transformOptional[Normalize], optional
Custom normalization transform. If None, uses default [-1,1] scaling.
Returns
- TorchImgNormalize
Normalized image in [-1,1] range with batch dimension (1, C, H, W)
- wibench.utils.denormalize_image(image: TorchImgNormalize, transform: Optional[Normalize] = None) TorchImg[source]
Denormalize image tensor from [-1,1] to [0,1] range and remove batch dimension.
Parameters
- imageTorchImgNormalize
Input image in [-1,1] range (1, C, H, W)
- transformOptional[Normalize], optional
Custom denormalization transform. If None, uses default scaling.
Returns
- TorchImg
Denormalized image in [0,1] range (C, H, W)
Implement the “extract” function
This function extracts a watermark from an attacked watermarked object. For example, for image-based algorithms, it takes the attacked image and watermark_data as input. It returns the extraction result, such as the extracted bit message.
from wibench.watermark_data import TorchBitWatermarkData
from wibench.typing import TorchImg
class MyAlgorithmWrapper(BaseAlgorithmWrapper):
...
def extract(self, image: TorchImg, watermark_data: TorchBitWatermarkData) -> Any:
...
Implemented algorithms
ARWGAN
- class wibench.algorithms.arwgan.wrapper.ARWGANParams(H: int, W: int, wm_length: int, encoder_blocks: int, encoder_channels: int, decoder_blocks: int, decoder_channels: int, use_discriminator: bool, use_vgg: bool, discriminator_blocks: int, discriminator_channels: int, decoder_loss: float, encoder_loss: float, adversarial_loss: float, enable_fp16: bool = False)[source]
Configuration parameters for the ARWGAN watermarking algorithm.
Attributes
- Hint
Height of the input image (in pixels). Determines the vertical size of image tensors
- Wint
Width of the input image (in pixels). Determines the horizontal size of image tensors
- wm_lengthint
Length of the binary watermark message to embed (in bits)
- encoder_blocksint
Number of convolutional blocks in the encoder network
- encoder_channelsint
Number of filters (channels) in each encoder block
- decoder_blocksint
Number of convolutional blocks in the decoder network
- decoder_channelsint
Number of filters in each decoder block
- use_discriminatorbool
If True, enables the use of an adversarial discriminator
- use_vggbool
If True, adds a perceptual loss using VGG features to improve
- discriminator_blocksint
Number of convolutional blocks in the discriminator network
- discriminator_channelsint
Number of filters in each discriminator block
- decoder_lossfloat
Weight of the decoder loss term in the total loss function. Controls the importance of accurate message recovery
- encoder_lossfloat
Weight of the encoder loss term in the total loss function. Typically regularizes visual similarity between original and encoded images
- adversarial_lossfloat
Weight of the adversarial loss term in the total loss. Higher values push the encoder to generate more realistic images when a discriminator is used
- enable_fp16bool
If True, enables mixed precision (fp16) training/inference for improved speed and reduced memory usage on compatible hardware (default False)
- class wibench.algorithms.arwgan.wrapper.ARWGANWrapper(*args, **kwargs)[source]
ARWGAN: Attention-Guided Robust Image Watermarking Model Based on GAN — Image Watermarking Algorithm.
Provides an interface for embedding and extracting watermarks using the ARWGAN watermarking algorithm. Based on the code from here.
Parameters
- paramsDict[str, Any]
ARWGAN algorithm configuration parameters (default EmptyDict)
- embed(image: TorchImg, watermark_data: TorchBitWatermarkData) TorchImg[source]
Embed watermark into input image.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: TorchBitWatermarkData
Torch bit message with data type torch.int64
CIN
- class wibench.algorithms.cin.wrapper.CINParams(H: int, W: int, wm_length: int, pre_noise_policy: PreNoisePolicy, experiment: str = '')[source]
Configuration parameters for the CIN algorithm.
Attributes
- Hint
Height of the input image (in pixels). Determines the vertical size of image tensors
- Wint
Width of the input image (in pixels). Determines the horizontal size of image tensors
- wm_lengthint
Length of the binary watermark message to embed (in bits)
- pre_noise_policyPreNoisePolicy
A policy that defines the parameters of noise for noise-specific selection module (NSM)
- experiment: str
The name of the experiment (default is “”)
- class wibench.algorithms.cin.wrapper.CINWrapper(*args, **kwargs)[source]
CIN: Towards Blind Watermarking: Combining Invertible and Non-invertible Mechanisms - Image Watermarking Algorithm [paper].
Provides an interface for embedding and extracting watermarks using the CIN watermarking algorithm. Based on the code from here.
Parameters
- paramsDict[str, Any]
CIN algorithm configuration parameters (default EmptyDict)
- embed(image: TorchImg, watermark_data: TorchBitWatermarkData) TorchImg[source]
Embed watermark into input image.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: TorchBitWatermarkData
Torch bit message with data type torch.int64
DCT
- class wibench.algorithms.dct_marker.wrapper.WatermarkData(watermark: ndarray, key: ndarray)[source]
Data for watermark.
- class wibench.algorithms.dct_marker.wrapper.DCTMarkerWrapper(params: Dict[str, Any] = {})[source]
Implementation of CAISS watermarking scheme via discrete cosine transform domain (“Correlation-and-Bit-Aware Spread Spectrum Embedding for Data Hiding”). Implementation is based on “Real data performance evaluation of CAISS watermarking scheme”.
Parameters
- paramsDict[str, Any]
dictionary, containing values for DCTMarkerConfig dataclass (default EmptyDict)
- embed(image: TorchImg, watermark_data: WatermarkData) TorchImg[source]
Embed watermark into input object (abstract).
- extract(image: TorchImg, watermark_data: WatermarkData) ndarray[source]
Extract watermark from marked object (abstract).
- watermark_data_gen() WatermarkData[source]
Generate watermark payload data.
Returns
- WatermarkData
Generated watermark payload
Notes
Default implementation returns None
Override to provide custom payload generation
Called automatically during embedding
DFT Circle
- class wibench.algorithms.dft_circle.wrapper.DFTMarkerWrapper(params: Dict[str, Any] = {})[source]
Implementation of image watermarking algorithm described in “Discrete Fourier transform-based watermarking method with an optimal implementation radius”.
Parameters
- paramsDict[str, Any]
Contains value for watermark strength “alpha” parameter of the algorithm (default EmptyDict)
DWSF
- class wibench.algorithms.dwsf.wrapper.DWSFParams(module_path: Union[str, pathlib.Path, NoneType] = None, device: str = 'cpu', encoder_weights_path: str = './model_files/dwsf/encoder_best.pth', decoder_weights_path: str = './model_files/dwsf/decoder_best.pth', seg_weights_path: str = './model_files/dwsf/seg.pth', message_length: int = 30, H: int = 128, W: int = 128, split_size: int = 128, default_noise_layer: List[str] = <factory>, mean: List[float] = <factory>, std: List[float] = <factory>, psnr: int = 35, gt: float = 0.5)[source]
- class wibench.algorithms.dwsf.wrapper.DWSFWrapper(*args, **kwargs)[source]
DWSF: Practical Deep Dispersed Watermarking with Synchronization and Fusion - Image Watermarking Algorithm.
Provides an interface for embedding and extracting watermarks using the DWSF watermarking algorithm. Based on the code from here.
Parameters
- paramsDict[str, Any]
DWSF algorithm configuration parameters (default EmptyDict)
- encode(images, messages, splitSize=128, inputSize=128, h_coor=[], w_coor=[], psnr=35)[source]
Encode image blocks based on random coordinates.
- embed(image: TorchImg, watermark_data: TorchBitWatermarkData) ndarray[source]
Embed watermark into input image.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: TorchBitWatermarkData
Torch bit message with data type torch.int64
DWT SVM
- class wibench.algorithms.dwt_svm.wrapper.WatermarkData(watermark: ndarray, key: ndarray)[source]
Watermark data for DWT_SVM watermarking algorithm.
- class wibench.algorithms.dwt_svm.wrapper.DWTSVMWrapper(params: dict[str, typing_extensions.Any] = {})[source]
Custom implementation of image watermarking algorithm described in the paper.
Parameters
- paramsDict[str, Any]
Contains value for “threshold” parameter of the algorithm. The higher is the threshold, the watermark is more robust to attacks, but less imperceptible (default EmptyDict)
- embed(image: TorchImg, watermark_data: WatermarkData) TorchImg[source]
Embed watermark into input object (abstract).
- extract(image: TorchImg, watermark_data: WatermarkData) ndarray[source]
Extract watermark from marked object (abstract).
- watermark_data_gen() WatermarkData[source]
Generate watermark payload data.
Returns
- WatermarkData
Generated watermark payload
Notes
Default implementation returns None
Override to provide custom payload generation
Called automatically during embedding
DWT DCT
- class wibench.algorithms.invisible_watermark.wrapper.DwtDctWrapper(params: Dict[str, Any] = {})[source]
Image watermarking using frequency-domain transforms: DWT + DCT.
Provides an interface for embedding and extracting watermarks using the frequency-domain transforms: DWT + DCT. Based on the code from here.
Parameters
- paramsDict[str, Any]
DWT-DCT algorithm configuration parameters
DWT DCT SVD
- class wibench.algorithms.invisible_watermark.wrapper.DwtDctSvdWrapper(params: Dict[str, Any] = {})[source]
Image frequency-domain watermarking with additional SVD processing: DWT + DCT + SVD.
Provides an interface for embedding and extracting watermarks using the frequency-domain with additional SVD processing: DWT + DCT + SVD. Based on the code from the github repository.
Parameters
- paramsDict[str, Any]
DWT-DCT-SVD algorithm configuration parameters
InvisMark
InvisMark: Invisible and Robust Watermarking for AI-generated Image Provenance
Provides an interface for embedding and extracting watermarks using the InvisMark watermarking algorithm. Based on the code from here.
Note: real capacity of InvisMark is 94 message bits (reffer to watermark_data_gen for more information)
MBRS
MBRS: Enhancing Robustness of DNN-based Watermarking by Mini-Batch of Real and Simulated JPEG Compression
Provides an interface for embedding and extracting watermarks using the MBRS watermarking algorithm. Based on the code from here.
RivaGAN
- class wibench.algorithms.invisible_watermark.wrapper.RivaGanWrapper(params: Dict[str, Any] = {})[source]
Image watermarking via RivaGAN: a deep-learning-based encoder/decoder with attention mechanism [repository].
Provides an interface for embedding and extracting watermarks using the RivaGAN watermarking algorithm. Based on the codes from here.
Parameters
- paramsDict[str, Any]
RivaGAN algorithm configuration parameters
SSL watermarking
- class wibench.algorithms.ssl_watermarking.wrapper.SSLMarkerWrapper(*args, **kwargs)[source]
Watermarking Images in Self-Supervised Latent-Spaces (SSL) — Image Watermarking Algorithm [paper].
Provides an interface for embedding and extracting watermarks using the SSL watermarking algorithm. Based on the code from here.
Parameters
- paramsDict[str, Any]
SSL algorithm configuration parameters (default EmptyDict)
- embed(image: TorchImg, watermark_data: Union[WatermarkZeroBitData, WatermarkMultiBitData]) TorchImg[source]
Embed watermark into input image.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: Union[Watermark0BitData, WatermarkMultiBitData]
Watermark data for SSL (Self-Supervised Learning) watermarking algorithm in multi-bit or zero-bit scenario
- extract(image: TorchImg, watermark_data: Union[WatermarkZeroBitData, WatermarkMultiBitData]) Union[Tensor, float][source]
Extract watermark from marked image.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: Union[Watermark0BitData, WatermarkMultiBitData]
Watermark data for SSL (Self-Supervised Learning) watermarking algorithm in multi-bit or zero-bit scenario
- watermark_data_gen() Union[WatermarkMultiBitData, WatermarkZeroBitData][source]
Generate watermark payload data for SLL (Self-Supervised Learning) watermarking algorithm in both multi-bit and zero-bit scenarios.
Returns
- Union[WatermarkMultiBitData, Watermark0BitData]
Watermark data for SSL (Self-Supervised Learning) watermarking algorithm in multi-bit or zero-bit scenario
Notes
Called automatically during embedding
Stable Signature
- class wibench.algorithms.stable_signature.wrapper.StableSignatureParams(module_path: Union[str, pathlib.Path, NoneType] = None, device: str = 'cpu', ldm_config_path: str = './submodules/stable_signature/v2-inference.yaml', ldm_checkpoint_path: str = './model_files/stable_signature/v2-1_512-ema-pruned.ckpt', ldm_decoder_path: str = './model_files/stable_signature/sd2_decoder.pth', decoder_path: str = './model_files/stable_signature/dec_48b_whit.torchscript.pt', model: str = 'WIBE-HuggingFace/stable-diffusion-2', secret: str = '111010110101000001010111010011010100010000100111')[source]
- class wibench.algorithms.stable_signature.wrapper.StableSignatureWrapper(*args, **kwargs)[source]
The Stable Signature: Rooting Watermarks in Latent Diffusion Models — Image Watermarking Algorithm [paper].
Provides an interface for embedding and extracting watermarks in Text2Image task using the StableSignature watermarking algorithm. Based on the code from here.
Parameters
- paramsDict[str, Any]
StableSignature algorithm configuration parameters (default EmptyDict)
- embed(prompt: str, watermark_data: TorchBitWatermarkData) TorchImg[source]
Embed watermark into input image.
Parameters
- promptstr
Input prompt for image generation
- watermark_data: TorchBitWatermarkData
Torch bit message with data type torch.int64
StegaStamp
- class wibench.algorithms.stega_stamp.wrapper.StegaStampParams(weights_path: str = './model_files/stega_stamp/stega_stamp.onnx', wm_length: int = 100, width: int = 400, height: int = 400, alpha: float = 1.0)[source]
- class wibench.algorithms.stega_stamp.wrapper.StegaStampWrapper(*args, **kwargs)[source]
StegaStamp: Invisible Hyperlinks in Physical Photographs — Image Watermarking Algorithm [paper].
Provides an interface for embedding and extracting watermarks using the StegaStamp watermarking algorithm. Based on the code from the github repository.
Parameters
- paramsDict[str, Any]
StegaStamp algorithm configuration parameters (default EmptyDict)
- embed(image: TorchImg, watermark_data: TorchBitWatermarkData) TorchImg[source]
Embed watermark into input image.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: TorchBitWatermarkData
Torch bit message with data type torch.int64
TreeRing
- class wibench.algorithms.treering.wrapper.TreeRingParams(module_path: Optional[Union[str, Path]] = None, device: str = 'cpu', run_name: str = 'test', dataset: str = 'Gustavosta/Stable-Diffusion-Prompts', start: int = 1, end: int = 10, image_length: int = 512, model_id: str = 'WIBE-HuggingFace/stable-diffusion-2-1-base', with_tracking: str = 'store_true', num_images: int = 1, guidance_scale: float = 7.5, num_inference_steps: int = 50, test_num_inference_steps: Optional[int] = None, reference_model: Optional[str] = None, reference_model_pretrain: Optional[str] = None, max_num_log_image: int = 100, gen_seed: int = 10, w_seed: int = 999999, w_channel: int = 0, w_pattern: str = 'rand', w_mask_shape: str = 'circle', w_radius: int = 10, w_measurement: str = 'l1_complex', w_injection: str = 'complex', w_pattern_const: int = 0, threshold: int = 77)[source]
Paramenters of Tree-ring watermarking algorithm.
- class wibench.algorithms.treering.wrapper.TreeRingWatermarkData(watermark: Tensor, watermarking_mask: Tensor, gt_patch: Tensor)[source]
Watermark data for Tree-ring watermarking algorithm.
Attributes
- watermarktorch.Tensor
Latent noise with embedded watermark
- watermarking_masktorch.Tensor
Watermarking noise pattern
- gt_patchtorch.Tensor
Ground-truth patch
- class wibench.algorithms.treering.wrapper.TreeRingWrapper(params: Dict[str, Any] = {})[source]
Tree-Ring Watermarks: Fingerprints for Diffusion Images that are Invisible and Robust - Image Watermarking Algorithm.
Provides an interface for embedding and extracting watermarks in Text2Image task using the Tree-Ring watermarking algorithm. Based on the code from here.
Parameters
- paramsDict[str, Any]
Tree-Ring algorithm configuration parameters (default EmptyDict)
- embed(prompt: str, watermark_data: TreeRingWatermarkData) TorchImg[source]
Generates a watermarked image based on a text prompt.
Parameters
- promptstr
Input prompt for image generation
- watermark_data: TreeRingWatermarkData
Watermark data for Tree-ring watermarking algorithm
- extract(img: TorchImg, watermark_data: TreeRingWatermarkData) bool[source]
Extract watermark from marked image.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: TreeRingWatermarkData
Watermark data for Tree-ring watermarking algorithm
Notes
Obtains latent values after DDIM inversion and compares them with a threshold
- watermark_data_gen() TreeRingWatermarkData[source]
Get watermark payload data for Tree-ring watermarking algorithm.
Returns
- TreeRingWatermarkData
Watermark data for Tree-ring watermarking algorithm
Notes
Called automatically during embedding
TrustMark
- class wibench.algorithms.trustmark.wrapper.TrustMarkParams(module_path: Optional[Union[str, Path]] = None, device: str = 'cpu', wm_length: int = 100, model_type: Literal['Q', 'B', 'C', 'P'] = 'Q', wm_strength: float = 0.75)[source]
Configuration parameters for the TrustMark algorithm.
Attributes
- wm_lengthint
Length of the watermark message to be embed (in bits) (default 100).
- model_typeLiteral[‘Q’, ‘B’, ‘C’]
Specifies the model architecture variant (default Q): - ‘Q’: (Quality) Trade-off between robustness and imperceptibility. Uses ResNet-50 decoder. - ‘B’: (Beta) Very similar to Q, included mainly for reproducing the paper. Uses ResNet-50 decoder. - ‘C’: (Compact). Uses a ResNet-18 decoder (smaller model size). Slightly lower visual quality. - ‘P’: (Perceptual). Very high visual quality and good robustness. ResNet-50 decoder trained with much higher weight on perceptual loss (see paper).
- wm_strengthfloat
Controls visibility/strength of watermark embedding (default 0.75)
- class wibench.algorithms.trustmark.wrapper.TrustMarkWrapper(*args, **kwargs)[source]
TrustMark: Universal Watermarking for Arbitrary Resolution Images - Image Watermarking Algorithm.
Provides an interface for embedding and extracting watermarks using the TrustMark watermarking algorithm. Based on the code from here.
Parameters
- paramsDict[str, Any]
TrustMark algorithm configuration parameters (default EmptyDict)
- embed(image: TorchImg, watermark_data: TorchBitWatermarkData)[source]
Embed watermark into input image.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: TorchBitWatermarkData
Torch bit message with data type torch.int64
VideoSeal, PixelSeal, ChunkySeal
- class wibench.algorithms.videoseal.wrapper.VideosealWrapper(*args, **kwargs)[source]
Video Seal: Open and Efficient Video Watermarking
Provides an interface for embedding and extracting watermarks using the VideoSeal watermarking algorithm. Based on the code from here.
- embed(image: TorchImg, watermark_data: TorchBitWatermarkData)[source]
Embed watermark into input object (abstract).
- class wibench.algorithms.videoseal.wrapper.PixelSeal(*args, **kwargs)[source]
Pixel Seal: Adversarial-only training for invisible image and video watermarking
Provides an interface for embedding and extracting watermarks using the PixelSeal watermarking algorithm. Based on the code from here.
- class wibench.algorithms.videoseal.wrapper.ChunkySeal(*args, **kwargs)[source]
We Can Hide More Bits: The Unused Watermarking Capacity in Theory and in Practice
Provides an interface for embedding and extracting watermarks using the ChunkySeal watermarking algorithm. Based on the code from here.
Note: Model weights are not provided by download_models.py script. You may get them from original link.
Watermark Anything
- class wibench.algorithms.watermark_anything.wrapper.WAParams(wm_length: int, scaling_w: float)[source]
Configuration parameters for the WA (Watermark Anything) watermarking algorithm.
Attributes
- wm_lengthint
Length of the watermark message to be embed (in bits).
- scaling_wfloat
Scaling factor for the watermark in the embedder model.
- class wibench.algorithms.watermark_anything.wrapper.WatermarkAnythingWrapper(*args, **kwargs)[source]
Watermark Anything with Localized Messages - Image Watermarking Algorithm [paper].
Provides an interface for embedding and extracting watermarks using the Watermark Anything watermarking algorithm. Based on the code from here.
- embed(image: TorchImg, watermark_data: TorchBitWatermarkData)[source]
Embed watermark into input image.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: TorchBitWatermarkData
Torch bit message with data type torch.int64
MaskWM
- class wibench.algorithms.maskwm.wrapper.WmEncoderConfig(message_length: int = 32, in_channels: int = 3, tp_channels: int = 3, mask_channel: int = 0, channels: int = 64, norm_type: str = 'group')[source]
- class wibench.algorithms.maskwm.wrapper.WmDecoderConfig(message_length: int = 32, in_channels: int = 3, tp_channels: int = 3, mask_channel: int = 1, channels: int = 128, norm_type: str = 'group')[source]
- class wibench.algorithms.maskwm.wrapper.MaskWMParams(module_path: Union[str, pathlib.Path, NoneType] = None, device: str = 'cpu', checkpoint_path: str = './model_files/maskwm/D_32bits.pth', use_jnd: bool = True, jnd_factor: float = 1.3, blue: bool = True, image_size: int = 256, wm_enc_config: wibench.algorithms.maskwm.wrapper.WmEncoderConfig = <factory>, wm_dec_config: wibench.algorithms.maskwm.wrapper.WmDecoderConfig = <factory>)[source]
- class wibench.algorithms.maskwm.wrapper.MaskWMWrapper(*args, **kwargs)[source]
Mask Image Watermarking — Image Watermarking Algorithm [paper].
Provides an interface for embedding and extracting watermarks using the MaskWM algorithm. Based on the code from here.
- embed(image: TorchImg, watermark_data: TorchBitWatermarkData) TorchImg[source]
Embed watermark into input image.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: TorchBitWatermarkData
Torch bit message with data type torch.int64
SyncSeal
- class wibench.algorithms.syncseal.wrapper.JNDConfig(in_channels: int = 1, out_channels: int = 1)[source]
- class wibench.algorithms.syncseal.wrapper.EmbedderConfig(model: str = 'unet_small2_yuv', in_channels: int = 1, out_channels: int = 1, z_channels: int = 16, num_blocks: int = 8, activation: str = 'gelu', normalization: str = 'group', z_channels_mults: List[int] = <factory>, last_tanh: bool = True)[source]
- class wibench.algorithms.syncseal.wrapper.ExtractorEncoder(depths: List[int] = <factory>, dims: List[int] = <factory>)[source]
- class wibench.algorithms.syncseal.wrapper.ExtractorHead(embed_dim: int = 768, out_dim: int = 8)[source]
- class wibench.algorithms.syncseal.wrapper.ExtractorConfig(model: str = 'convnext_tiny', encoder: Dict[str, Any] = <factory>, head: Dict[str, Any] = <factory>)[source]
- class wibench.algorithms.syncseal.wrapper.SyncSealParams(module_path: Union[str, pathlib.Path, NoneType] = None, device: str = 'cpu', checkpoint_path: str = './model_files/syncseal/syncmodel.jit.pt', img_size_proc: int = 256, scaling_i: float = 1.0, scaling_w: float = 0.2, embedder_config: wibench.algorithms.syncseal.wrapper.EmbedderConfig = <factory>, extractor_config: wibench.algorithms.syncseal.wrapper.ExtractorConfig = <factory>, jnd_config: wibench.algorithms.syncseal.wrapper.JNDConfig = <factory>, method: str = 'trustmark', method_params: Dict[str, Any] = <factory>)[source]
- class wibench.algorithms.syncseal.wrapper.SyncSeal(*args, **kwargs)[source]
GEOMETRIC IMAGE SYNCHRONIZATION WITH DEEP WATERMARKING — Image Synchronization Algorithm [paper].
Provides an interface for embedding and extracting watermarks using the SyncSeal synchronization algorithm with selected image watermarking algorithm. Based on the code from the github repository.
Parameters
- paramsDict[str, Any]
SyncSeal algorithm configuration parameters (default EmptyDict)
- embed(image: TorchImg, watermark_data: Any) TorchImg[source]
Embed both watermarking, marking methods and synchronization.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: TorchBitWatermarkData
Torch bit message with data type torch.int64
Gaussian Shading
Ring-ID
- class wibench.algorithms.ringid.wrapper.RingIDParams(module_path: ~typing.Optional[~typing.Union[str, ~pathlib.Path]] = None, device: str = 'cpu', radius: int = 14, radius_cutoff: int = 3, anchor_x_offset: int = 0, anchor_y_offset: int = 0, use_rounder_ring: bool = True, ring_value_range: int = 64, quantization_levels: int = 2, assigned_keys: int = -1, fix_gt: int = 1, time_shift: int = 1, heter_watermark_channel: ~typing.List[int] = <factory>, ring_watermark_channel: ~typing.List[int] = <factory>, mode: str = 'complex', p: int = 1, channel_min: int = 1, image_length: int = 512, model_id: str = 'WIBE-HuggingFace/stable-diffusion-2-1-base', with_tracking: str = 'store_true', num_images: int = 1, guidance_scale: float = 7.5, num_inference_steps: int = 50, test_num_inference_steps: ~typing.Optional[int] = None, threshold: float = 50)[source]
Paramenters of RingID watermarking algorithm.
- class wibench.algorithms.ringid.wrapper.RignIDWatermarkData(watermark_pattern: Tensor, watermark_mask: Tensor)[source]
Watermark data for RingID watermarking algorithm.
Attributes
- watermark_patterntorch.Tensor
Latent noise with embedded watermark
- watermark_masktorch.Tensor
Watermarking noise pattern
- class wibench.algorithms.ringid.wrapper.RingIDWrapper(params: Dict[str, Any] = {})[source]
RingID: Rethinking Tree-Ring Watermarking for Enhanced Multi-Key Identification - Image Watermarking Algorithm.
Provides an interface for embedding and extracting watermarks in Text2Image task using the RingID watermarking algorithm. Based on the code from here.
Parameters
- paramsDict[str, Any]
RingID algorithm configuration parameters (default EmptyDict)
- embed(prompt: str, watermark_data: RignIDWatermarkData) TorchImg[source]
Generates a watermarked image based on a text prompt.
Parameters
- promptstr
Input prompt for image generation
- watermark_data: RingIDWatermarkData
Watermark data for RingID watermarking algorithm
- extract(img: TorchImg, watermark_data: RignIDWatermarkData) bool[source]
Extract watermark from marked image.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: RingIDWatermarkData
Watermark data for RingID watermarking algorithm
Notes
Obtains latent values after DDIM inversion and compares them with a threshold
- watermark_data_gen() RignIDWatermarkData[source]
Get watermark payload data for RingID watermarking algorithm.
Returns
- RingIDWatermarkData
Watermark data for RingID watermarking algorithm
Notes
Called automatically during embedding
MaXsive
- class wibench.algorithms.maxsive.wrapper.MaXsiveParams(module_path: Optional[Union[str, Path]] = None, device: str = 'cpu', model_id: str = 'WIBE-HuggingFace/stable-diffusion-2-1-base', model_name: str = 'watermarkSD21', guidance_scale: float = 7.5, num_inference_steps: int = 50, num_inversion_steps: Optional[int] = None, num_images_per_prompt: int = 1, channel_copy: int = 1, hw_copy: int = 2, template_c: int = 3, distant_func: str = 'corr', diffusion_bit: int = 16, tpr_file: str = './submodules/MaXsive/threshold/MaXsive-cos.pt')[source]
Paramenters of MaXsive watermarking algorithm.
- class wibench.algorithms.maxsive.wrapper.MaXsiveWatermarkData(watermark: Tensor, z: Tensor, data: Dict[str, Any])[source]
Watermark data for RingID watermarking algorithm.
Attributes
- watermarktorch.Tensor
Normalized sequence of normally distributed numbers
- ztorch.Tensor
Latent noise with embedded watermark
- dataDict[str, Any]
Data for watermark extraction
- class wibench.algorithms.maxsive.wrapper.MaXsiveWrapper(params: Dict[str, Any] = {})[source]
MaXsive: High-Capacity and Robust Training-Free Generative Image Watermarking in Diffusion Models.
Provides an interface for embedding and extracting watermarks in Text2Image task using the MaXsive watermarking algorithm. Based on the code from here.
Parameters
- paramsDict[str, Any]
MaXsive algorithm configuration parameters (default EmptyDict)
- embed(prompt: str, watermark_data: MaXsiveWatermarkData) TorchImg[source]
Generates a watermarked image based on a text prompt.
Parameters
- promptstr
Input prompt for image generation
- watermark_data: MaXsiveWatermarkData
Watermark data for MaXsive watermarking algorithm
- extract(img: TorchImg, watermark_data: MaXsiveWatermarkData) bool[source]
Extract watermark from marked image.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: MaXsiveWatermarkData
Watermark data for MaXsive watermarking algorithm
- watermark_data_gen() MaXsiveWatermarkData[source]
Get watermark payload data for MaXsive watermarking algorithm.
Returns
- MaXsiveWatermarkData
Watermark data for MaXsive watermarking algorithm
Notes
Called automatically during embedding
METR
- class wibench.algorithms.metr.wrapper.METRParams(module_path: Optional[Union[str, Path]] = None, device: str = 'cpu', model_id: str = 'WIBE-HuggingFace/stable-diffusion-2-1-base', model_name: str = 'watermarkSD21', guidance_scale: float = 7.5, num_inference_steps: int = 40, num_inversion_steps: Optional[int] = None, num_images_per_prompt: int = 1, channel_copy: int = 1, hw_copy: int = 2, template_c: int = 3, distant_func: str = 'corr', diffusion_bit: int = 16, tpr_file: Optional[str] = None, image_length: int = 512, w_radius: int = 10, w_seed: int = 999999, w_pattern: str = 'ring', w_pattern_const: float = 0.0, w_injection: str = 'complex', w_channel: int = 3, w_mask_shape: str = 'circle', use_random_msgs: bool = True, msg_type: str = 'binary', msg: Optional[str] = None, msg_scaler: int = 100, w_measurement: str = 'l1_complex', decoder_state_dict_path: Optional[str] = None, stable_sig_full_model_config: Optional[str] = None, stable_sig_full_model_ckpt: Optional[str] = None)[source]
Paramenters of METR watermarking algorithm.
- class wibench.algorithms.metr.wrapper.METRWatermarkData(watermark: Tensor, watermark_mask: Tensor, init_latents_w: Tensor)[source]
Watermark data for METR watermarking algorithm.
Attributes
- watermarktorch.Tensor
Torch bit message
- watermark_masktorch.Tensor
Watermark noise pattern
- init_latents_wDict[str, Any]
Latent noise with embedded watermark
- class wibench.algorithms.metr.wrapper.METRWrapper(params: Dict[str, Any] = {})[source]
METR: Image Watermarking with Large Number of Unique Messages.
Provides an interface for embedding and extracting watermarks in Text2Image task using the METR watermarking algorithm. Based on the code from here.
Parameters
- paramsDict[str, Any]
METR algorithm configuration parameters (default EmptyDict)
- embed(prompt: str, watermark_data: METRWatermarkData) TorchImg[source]
Generates a watermarked image based on a text prompt.
Parameters
- promptstr
Input prompt for image generation
- watermark_data: METRWatermarkData
Watermark data for METR watermarking algorithm
- extract(img: TorchImg, watermark_data: METRWatermarkData) bool[source]
Extract watermark from marked image.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: METRWatermarkData
Watermark data for METR watermarking algorithm
- watermark_data_gen() METRWatermarkData[source]
Get watermark payload data for METR watermarking algorithm.
Returns
- METRWatermarkData
Watermark data for METR watermarking algorithm
Notes
Called automatically during embedding
PIMoG
- class wibench.algorithms.pimog.wrapper.PIMoGParams(module_path: Union[str, pathlib.Path, NoneType] = None, device: str = 'cpu', checkpoint_path: str = './model_files/pimog/Encoder_Decoder_Model_mask_99.pth', image_size: int = 128, wm_length: int = 30)[source]
- class wibench.algorithms.pimog.wrapper.PIMoGWrapper(*args, **kwargs)[source]
PIMoG: An Effective Screen-shooting Noise-Layer Simulation for Deep-Learning-Based Watermarking Network — Image Watermarking Algorithm [paper].
Provides an interface for embedding and extracting watermarks using the PIMoG watermarking algorithm. Based on the code from the github repository.
Parameters
- paramsDict[str, Any]
PIMoG algorithm configuration parameters (default EmptyDict)
- embed(image: TorchImg, watermark_data: TorchBitWatermarkData) TorchImg[source]
Embed watermark into input image.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: TorchBitWatermarkData
Torch bit message with data type torch.int64
Robust-Wide
- class wibench.algorithms.robust_wide.wrapper.RobustWideEncoderParams(image_size: int = 512, message_length: int = 64, in_channels: int = 3, channels: int = 64, norm_type: str = 'batch', final_skip: bool = True)[source]
- class wibench.algorithms.robust_wide.wrapper.RobustWideDecoderParams(image_size: int = 512, message_length: int = 64, in_channels: int = 3, norm_type: str = 'batch')[source]
- class wibench.algorithms.robust_wide.wrapper.RobustWideWmModelParams(wm_enc_config: wibench.algorithms.robust_wide.wrapper.RobustWideEncoderParams = <factory>, wm_dec_config: wibench.algorithms.robust_wide.wrapper.RobustWideDecoderParams = <factory>)[source]
- class wibench.algorithms.robust_wide.wrapper.RobustWideParams(module_path: Union[str, pathlib.Path, NoneType] = None, device: str = 'cpu', checkpoint_path: str = './model_files/robust_wide/wm_model.ckpt', wm_model_config: wibench.algorithms.robust_wide.wrapper.RobustWideWmModelParams = <factory>)[source]
- class wibench.algorithms.robust_wide.wrapper.RobustWideWrapper(*args, **kwargs)[source]
Robust-Wide: Robust Watermarking Against Instruction-Driven Image Editing — Image Watermarking Algorithm [paper].
Provides an interface for embedding and extracting watermarks using the Robust-Wide watermarking algorithm. Based on the code from the github repository.
Parameters
- paramsDict[str, Any]
Robust-Wide algorithm configuration parameters (default EmptyDict)
- embed(image: TorchImg, watermark_data: TorchBitWatermarkData) TorchImg[source]
Embed watermark into input image.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: TorchBitWatermarkData
Torch bit message with data type torch.int64
FIN
- class wibench.algorithms.fin.wrapper.FINParams(module_path: Union[str, pathlib.Path, NoneType] = None, device: str = 'cpu', H: int = 128, W: int = 128, wm_length: int = 64, fed_checkpoint: str = './model_files/fin/jpeg/FED.pt')[source]
- class wibench.algorithms.fin.wrapper.FINWrapper(*args, **kwargs)[source]
FIN: Flow-Based Robust Watermarking with Invertible Noise Layer for Black-Box Distortions — Image Watermarking Algorithm [paper].
Provides an interface for embedding and extracting watermarks using the FIN watermarking algorithm. Based on the code from here here.
Parameters
- paramsDict[str, Any]
FIN algorithm configuration parameters (default EmptyDict)
- embed(image: TorchImg, watermark_data: TorchBitWatermarkData) TorchImg[source]
Embed watermark into input image.
Parameters
- imageTorchImg
Input image tensor in (C, H, W) format
- watermark_data: TorchBitWatermarkData
Torch bit message with data type torch.int64