# -*- coding: utf-8 -*-
import torch
from torch import nn
[docs]class BasicDecoder(nn.Module):
"""
The BasicDecoder module takes an steganographic image and attempts to decode
the embedded data tensor.
Input: (N, 3, H, W)
Output: (N, D, H, W)
"""
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.layers = nn.Sequential(
self._conv2d(3, 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, self.hidden_size),
nn.LeakyReLU(inplace=True),
nn.BatchNorm2d(self.hidden_size),
self._conv2d(self.hidden_size, self.data_depth)
)
return [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._models = [self.layers]
self.version = '1'
[docs] def forward(self, x):
x = self._models[0](x)
if len(self._models) > 1:
x_list = [x]
for layer in self._models[1:]:
x = layer(torch.cat(x_list, dim=1))
x_list.append(x)
return x
[docs]class DenseDecoder(BasicDecoder):
"""
The DenseDecoder module takes an steganographic image and attempts to decode
the embedded data tensor.
Input: (N, 3, H, W)
Output: (N, D, H, W)
"""
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.hidden_size),
nn.LeakyReLU(inplace=True),
nn.BatchNorm2d(self.hidden_size)
)
self.conv3 = nn.Sequential(
self._conv2d(self.hidden_size * 2, 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))
return self.conv1, self.conv2, self.conv3, self.conv4
[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._models = [
self.conv1,
self.conv2,
self.conv3,
self.conv4
]
self.version = '1'