| --- |
| license: apache-2.0 |
| tags: |
| - steganography |
| - watermarking |
| - image-watermarking |
| - pytorch |
| --- |
| |
| # PicoTrust — Robust Image Steganography |
|
|
| PicoTrust is a neural image steganography model that encodes hidden binary messages into images, surviving real-world distortions like JPEG compression, blur, noise, and color shifts. |
|
|
| Based on StegaStamp (Tancik et al., CVPR 2020) with significant architectural improvements. |
|
|
| ## Architecture |
|
|
| - **Encoder**: U-Net (512x512) + E_post refinement (1-channel grayscale residual, softsign bounding with strength annealing) |
| - **Decoder**: StegaStamp-style CNN (256x256) with compact STN for geometric alignment |
| - **Discriminator**: WGAN PatchGAN |
| - **Message**: 100 bits per image |
| - **Bounding**: Softsign residual with strength annealing (1.0 → target over training) |
| |
| ## Checkpoints |
| |
| | Model | File | Strength | PSNR | Bit Accuracy | JPEG Q10 | |
| |-------|------|----------|------|-------------|----------| |
| | **v4** | `v4/picotrust_v4_200k.pt` | 0.02 | **35.56 dB** | 97.8% | 97.6% | |
| | **v2** | `v2/picotrust_v2_200k.pt` | 0.03 | 32.82 dB | **98.4%** | **98.6%** | |
| |
| - **v4**: Best balance of visual quality and accuracy. +2.7 dB PSNR over v2 with <1% accuracy trade-off. |
| - **v2**: Highest raw accuracy and JPEG robustness. Zero colour shifts (architectural guarantee). |
| |
| Both models have zero colour shifts (grayscale residual) and >92% accuracy across all distortion types. |
| |
| ## Robustness (v4 @ 200k steps) |
| |
| | Distortion | Bit Accuracy | |
| |------------|-------------| |
| | Clean | 97.8% | |
| | JPEG Q10 | 97.6% | |
| | Gaussian blur (σ=3) | 98.0% | |
| | Gaussian noise (σ=0.05) | 96.6% | |
| | Brightness ±0.3 | 94.6% | |
| | Contrast ±0.3 | 98.2% | |
| |
| ## Usage |
| |
| ```python |
| import torch |
| from picode.models.picotrust import Encoder, Decoder |
| |
| # Load checkpoint |
| ckpt = torch.load("picotrust_v4_200k.pt", map_location="cpu") |
|
|
| # Reconstruct encoder/decoder |
| encoder = Encoder(message_length=100, image_size=512) |
| decoder = Decoder(message_length=100, image_size=256) |
|
|
| encoder.load_state_dict(ckpt["encoder"]) |
| decoder.load_state_dict(ckpt["decoder"]) |
|
|
| # Encode: image (1,3,512,512) [0,1] + message (1,100) binary |
| encoded = encoder(image, message) |
|
|
| # Decode: returns probabilities (1,100) in [0,1] |
| decoded = decoder(encoded) |
| bits = (decoded > 0.5).float() |
| ``` |
| |
| ## Training |
| |
| Trained on COCO train2017 (~118K images) for 200K steps on a single NVIDIA T4 GPU. |
| |
| Config: `configs/picotrust_v4.yaml` (v4) / `configs/picotrust_v2.yaml` (v2) |
| |
| ## Key Design Decisions |
| |
| 1. **Grayscale residual**: 1-channel E_post output broadcast to 3 channels — eliminates colour shifts by construction |
| 2. **Softsign bounding**: `strength * x / (1 + |x|)` — non-vanishing gradients unlike tanh |
| 3. **Strength annealing**: Start unbounded (1.0) → anneal to target — prevents bootstrap collapse |
| 4. **MSE message loss**: No trivial equilibrium unlike BCE |
| 5. **Zero-init E_post**: Residual starts at exactly zero, grows gradually |
| |
| ## License |
| |
| Apache 2.0 |
| |
