# -*- coding: utf-8 -*-
import torch
from torch import nn
[docs]class BasicEncoder(nn.Module):
"""
The BasicEncoder module takes an cover image and a data tensor and combines
them into a steganographic image.
Input: (N, 3, H, W), (N, D, H, W)
Output: (N, 3, H, W)
"""
add_image = False
def _conv2d(self, in_channels, out_channels):
return nn.Conv2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=3,
padding=1
)
def _build_models(self):
self.features = nn.Sequential(
self._conv2d(3, self.hidden_size),
nn.LeakyReLU(inplace=True),
nn.BatchNorm2d(self.hidden_size),
)
self.layers = nn.Sequential(
self._conv2d(self.hidden_size + self.data_depth, self.hidden_size),
nn.LeakyReLU(inplace=True),
nn.BatchNorm2d(self.hidden_size),
self._conv2d(self.hidden_size, self.hidden_size),
nn.LeakyReLU(inplace=True),
nn.BatchNorm2d(self.hidden_size),
self._conv2d(self.hidden_size, 3),
nn.Tanh(),
)
return self.features, self.layers
def __init__(self, data_depth, hidden_size):
super().__init__()
self.version = '1'
self.data_depth = data_depth
self.hidden_size = hidden_size
self._models = self._build_models()
[docs] def upgrade_legacy(self):
"""Transform legacy pretrained models to make them usable with new code versions."""
# Transform to version 1
if not hasattr(self, 'version'):
self.version = '1'
[docs] def forward(self, image, data):
x = self._models[0](image)
x_list = [x]
for layer in self._models[1:]:
x = layer(torch.cat(x_list + [data], dim=1))
x_list.append(x)
if self.add_image:
x = image + x
return x
[docs]class ResidualEncoder(BasicEncoder):
"""
The ResidualEncoder module takes an cover image and a data tensor and combines
them into a steganographic image.
Input: (N, 3, H, W), (N, D, H, W)
Output: (N, 3, H, W)
"""
add_image = True
def _build_models(self):
self.features = nn.Sequential(
self._conv2d(3, self.hidden_size),
nn.LeakyReLU(inplace=True),
nn.BatchNorm2d(self.hidden_size),
)
self.layers = nn.Sequential(
self._conv2d(self.hidden_size + self.data_depth, self.hidden_size),
nn.LeakyReLU(inplace=True),
nn.BatchNorm2d(self.hidden_size),
self._conv2d(self.hidden_size, self.hidden_size),
nn.LeakyReLU(inplace=True),
nn.BatchNorm2d(self.hidden_size),
self._conv2d(self.hidden_size, 3),
)
return self.features, self.layers
[docs]class DenseEncoder(BasicEncoder):
"""
The DenseEncoder module takes an cover image and a data tensor and combines
them into a steganographic image.
Input: (N, 3, H, W), (N, D, H, W)
Output: (N, 3, H, W)
"""
add_image = True
def _build_models(self):
self.conv1 = nn.Sequential(
self._conv2d(3, self.hidden_size),
nn.LeakyReLU(inplace=True),
nn.BatchNorm2d(self.hidden_size),
)
self.conv2 = nn.Sequential(
self._conv2d(self.hidden_size + self.data_depth, self.hidden_size),
nn.LeakyReLU(inplace=True),
nn.BatchNorm2d(self.hidden_size),
)
self.conv3 = nn.Sequential(
self._conv2d(self.hidden_size * 2 + self.data_depth, self.hidden_size),
nn.LeakyReLU(inplace=True),
nn.BatchNorm2d(self.hidden_size),
)
self.conv4 = nn.Sequential(
self._conv2d(self.hidden_size * 3 + self.data_depth, 3)
)
return self.conv1, self.conv2, self.conv3, self.conv4