-
Notifications
You must be signed in to change notification settings - Fork 7
/
Copy pathdecoder.py
91 lines (78 loc) · 2.71 KB
/
decoder.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
import torch
from torch import nn
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.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, self.hidden_size),
nn.LeakyReLU(inplace=True),
nn.BatchNorm2d(self.hidden_size),
)
self.conv4 = nn.Sequential(
self._conv2d(self.hidden_size, self.data_depth)
)
return self.conv1, self.conv2, self.conv3, self.conv4
def forward(self, image):
x = self._models[0](image)
x_1 = self._models[1](x)
x_2 = self._models[2](x_1)
x_3 = self._models[3](x_2)
return x_3
def __init__(self, data_depth, hidden_size):
super().__init__()
self.data_depth = data_depth
self.hidden_size = hidden_size
self._models = self._build_models()
class DenseDecoder(BasicDecoder):
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
def forward(self, image):
x = self._models[0](image)
x_list = [x]
x_1 = self._models[1](torch.cat(x_list, dim=1))
x_list.append(x_1)
x_2 = self._models[2](torch.cat(x_list, dim=1))
x_list.append(x_2)
x_3 = self._models[3](torch.cat(x_list, dim=1))
x_list.append(x_3)
return x_3