Robust and Imperceptible Black-Box DNN Watermarking (Paper Reproduction)

This repository is a reproduction implementation of the paper:

Robust and Imperceptible Black-Box DNN Watermarking Based on Fourier Perturbation Analysis and Frequency Sensitivity Clustering
Paper: https://arxiv.org/abs/2208.03944

The project reproduces the core idea of embedding ownership watermarks in the frequency domain and verifying model ownership in a black-box setting.

Reproduction Scope

This codebase focuses on reproducing the paper's main pipeline:

1. Train a clean model on Tiny-ImageNet (M0)
2. Perform Fourier perturbation analysis to compute frequency sensitivity
3. Use K-means clustering to generate a binary frequency mask (clustering_map)
4. Inject watermark triggers in Fourier domain on poisoned samples
5. Train watermarked model and evaluate with TA / WSR / image quality metrics

Key Features

• Fourier-domain watermark embedding (fft2 -> mask perturbation -> ifft2)
• Frequency sensitivity analysis and clustering-based trigger selection
• End-to-end pipeline for training + verification
• Support for multiple backbones from CLI (resnet34, googlenet)
• Auto-generated visual outputs (clean vs watermarked, FFT spectrum, trigger samples)

Project Structure

• main.py: full pipeline entry (train clean -> analyze -> train watermark -> verify)
• fourier_analysis.py: Fourier perturbation analysis + frequency clustering
• watermark.py: watermark embedding in frequency domain
• train_clean.py: clean model training
• train.py: poisoned training (watermark + label flip)
• verify.py: TA / WSR / PSNR / SSIM evaluation + visualization
• models.py: backbone definitions and model factory
• dataset.py: Tiny-ImageNet loading and preprocessing

Environment

Python 3.8+ is recommended.

Install dependencies:

pip install -r requirements.txt

Quick Start

1) Run full reproduction pipeline

python main.py --model resnet34 --save-dir output_resnet

or:

python main.py --model googlenet --save-dir output_googlenet

2) Common arguments

• --model {resnet34,googlenet}: choose architecture without changing code
• --save-dir: output root directory
• --perturbation-strength: watermark strength in frequency domain
• --poison-ratio: poisoned sample ratio
• --target-class: target label for watermark trigger
• --force-retrain: retrain even if checkpoints already exist
• --force-analysis: recompute Fourier analysis and clustering map

Outputs

Under --save-dir, the pipeline generates:

• models/clean_<model>.pth
• models/watermarked_<model>.pth
• images/clustering_map.npy
• images/clean_vs_watermarked_samples.png
• images/fft_spectrum_comparison.png
• images/trigger_samples.png
• logs/train_clean.log
• logs/fourier_analysis.log
• logs/train_watermark.log
• logs/verify.log

Evaluation Metrics

• TA (Task Accuracy): clean validation accuracy
• WSR (Watermark Success Rate): trigger success rate to target class
• PSNR / SSIM: perceptual similarity between clean and watermarked images

Notes on Reproduction Fidelity

As a reproduction project, results can vary with:

• model backbone choice (resnet34 vs googlenet)
• watermark strength and poison ratio
• training schedule and random seed
• hardware / CUDA / dependency versions

This repository aims to reproduce the core methodology and behavior trend described in the original paper, rather than claiming exact numeric parity under all environments.

Citation

If you use this reproduction code, please cite the original paper:

Robust and Imperceptible Black-Box DNN Watermarking Based on Fourier Perturbation Analysis and Frequency Sensitivity Clustering
arXiv: https://arxiv.org/abs/2208.03944