README.md

---
language:
- en
tags:
- image-watermarking
- onnx
- computer-vision
---

# Watermark Anything Model (WAM) - ONNX

This repository contains the ONNX conversion of the **Watermark Anything Model (WAM)**. To ensure optimal performance and flexibility, the monolithic PyTorch model has been decomposed into three distinct ONNX computational graphs.

## Model Components

The watermarking pipeline is divided into three `.onnx` files, ensuring that image scaling and heavy computations (like Just Noticeable Difference - JND attenuation) can be performed optimally.

1. **`embedder.onnx`**
   - **Inputs:**
     - `image`: Float32 `[batch_size, 3, 256, 256]` (Normalized image)
     - `message`: Float32 `[batch_size, nbits]` (Binary payload to embed)
   - **Outputs:**
     - `watermark`: Float32 `[batch_size, 3, 256, 256]` (Watermark signal/deltas)

2. **`blender.onnx`**
   - **Inputs:** 
     - `imgs`: Float32 `[batch_size, 3, height, width]` (Original normalized image at any resolution)
     - `deltas`: Float32 `[batch_size, 3, height, width]` (Watermark signal resized to original image dimensions)
   - **Outputs:**
     - `watermarked_image`: Float32 `[batch_size, 3, height, width]` (Final normalized image with JND applied)

3. **`extractor.onnx`**
   - **Inputs:**
     - `image`: Float32 `[batch_size, 3, 256, 256]` (Watermarked image, resized and normalized)
   - **Outputs:**
     - `mask`: Float32 `[batch_size, 1 + nbits, 256, 256]` (Detected watermark mask and extracted message bits)

---

## Inference Guide (Python)

You will need `onnxruntime`, `numpy`, `torch`, `torchvision`, and `Pillow`.

### 1. Installation

```bash
pip install onnxruntime numpy torch torchvision pillow
```

### 2. End-to-End Example

Below is a complete Python script to load an image, embed a watermark, save the result, and extract the watermark mask.

```python
import numpy as np
import onnxruntime as ort
import torch
from torchvision import transforms
from PIL import Image

# ---------------------------------------------------------
# 1. Initialization
# ---------------------------------------------------------
# Load ONNX Runtime sessions
embedder_sess = ort.InferenceSession("embedder.onnx")
blender_sess = ort.InferenceSession("blender.onnx")
extractor_sess = ort.InferenceSession("extractor.onnx")

# Define ImageNet normalization stats
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]

normalize = transforms.Normalize(mean=mean, std=std)
unnormalize = transforms.Normalize(
    mean=[-m/s for m, s in zip(mean, std)], 
    std=[1/s for s in std]
)
resize_256 = transforms.Resize((256, 256), antialias=True)
to_tensor = transforms.ToTensor()

# Configure watermark bits (0 means only localizer mask, no payload)
nbits = 0 

# ---------------------------------------------------------
# 2. Image Preprocessing
# ---------------------------------------------------------
img_path = "input_image.jpg"
pil_img = Image.open(img_path).convert("RGB")
img_tensor = to_tensor(pil_img).unsqueeze(0)  # [1, 3, H, W]

# Normalize for the model
img_normalized = normalize(img_tensor)
original_h, original_w = img_tensor.shape[2], img_tensor.shape[3]

# ---------------------------------------------------------
# 3. Embedding the Watermark Signal
# ---------------------------------------------------------
img_resized = resize_256(img_normalized)
msg = np.random.randint(0, 2, (1, nbits)).astype(np.float32)

embedder_inputs = {
    'image': img_resized.numpy(),
    'message': msg
}
# Output is [1, 3, 256, 256]
watermark_deltas = embedder_sess.run(None, embedder_inputs)[0] 

# ---------------------------------------------------------
# 4. Blending & JND Attenuation
# ---------------------------------------------------------
# Resize the watermark signal back to the original image resolution
inverse_resize = transforms.Resize((original_h, original_w), antialias=True)
watermark_deltas_full = inverse_resize(torch.from_numpy(watermark_deltas)).numpy()

blender_inputs = {
    'imgs': img_normalized.numpy(),
    'deltas': watermark_deltas_full
}
# Output is [1, 3, H, W] normalized
watermarked_img_normalized = blender_sess.run(None, blender_inputs)[0]

# ---------------------------------------------------------
# 5. Postprocessing and Saving
# ---------------------------------------------------------
watermarked_img_tensor = torch.from_numpy(watermarked_img_normalized)
watermarked_img_unnorm = unnormalize(watermarked_img_tensor)
watermarked_img_clamped = torch.clamp(watermarked_img_unnorm, 0, 1)

output_img = transforms.ToPILImage()(watermarked_img_clamped.squeeze(0))
output_img.save("watermarked_image.png")
print("✅ Watermarked image saved!")

# ---------------------------------------------------------
# 6. Extraction Detection
# ---------------------------------------------------------
# Resize the watermarked image to 256x256 before feeding into the extractor
extractor_input = resize_256(watermarked_img_tensor).numpy()

extractor_inputs = {
    'image': extractor_input
}
extracted_mask = extractor_sess.run(None, extractor_inputs)[0]

print(f"✅ Extracted mask shape: {extracted_mask.shape}")
```

## Performance & Optimization
Replacing the native monolithic implementation with these ONNX models offers massive portability, meaning it can be run seamlessly on edge devices using ONNX Runtime (CPU, CUDA, TensorRT, CoreML).

*Note: The `blender.onnx` relies on **dynamic axes** for image resolution, allowing you to avoid resizing your pristine original images. The `embedder` and `extractor` graphs operate strictly at `256x256` to align with their training resolutions.*


---
license: mit
---