| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class DoubleConv(nn.Module): |
| def __init__(self, cin, cout): |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.Conv2d(cin, cout, 3, padding=1), nn.BatchNorm2d(cout), nn.ReLU(inplace=True), |
| nn.Conv2d(cout, cout, 3, padding=1), nn.BatchNorm2d(cout), nn.ReLU(inplace=True), |
| ) |
| def forward(self, x): return self.net(x) |
|
|
|
|
|
|
| class UNet(nn.Module): |
| def __init__(self, in_ch=3, out_ch=3, base=48): |
| super().__init__() |
| self.d1 = DoubleConv(in_ch, base) |
| self.d2 = DoubleConv(base, base * 2) |
| self.d3 = DoubleConv(base * 2, base * 4) |
| self.pool = nn.MaxPool2d(2) |
| self.bottleneck = DoubleConv(base * 4, base * 8) |
| self.up3 = nn.ConvTranspose2d(base * 8, base * 4, 2, stride=2) |
| self.u3 = DoubleConv(base * 8, base * 4) |
| self.up2 = nn.ConvTranspose2d(base * 4, base * 2, 2, stride=2) |
| self.u2 = DoubleConv(base * 4, base * 2) |
| self.up1 = nn.ConvTranspose2d(base * 2, base, 2, stride=2) |
| self.u1 = DoubleConv(base * 2, base) |
| self.out = nn.Conv2d(base, out_ch, 1) |
|
|
| def forward(self, x): |
| c1 = self.d1(x) |
| c2 = self.d2(self.pool(c1)) |
| c3 = self.d3(self.pool(c2)) |
| b = self.bottleneck(self.pool(c3)) |
| x = self.u3(torch.cat([self.up3(b), c3], 1)) |
| x = self.u2(torch.cat([self.up2(x), c2], 1)) |
| x = self.u1(torch.cat([self.up1(x), c1], 1)) |
| return self.out(x) |
|
|
|
|
|
|
| class StegEncoder(nn.Module): |
| def __init__(self, msg_bits=64): |
| super().__init__() |
| self.msg_bits = msg_bits |
| self.unet = UNet(in_ch=3 + msg_bits, out_ch=3, base=48) |
|
|
| def forward(self, cover, msg): |
| B, _, H, W = cover.shape |
| msg_plane = msg.view(B, self.msg_bits, 1, 1).expand(B, self.msg_bits, H, W) |
| x = torch.cat([cover, msg_plane], dim=1) |
| residual = torch.tanh(self.unet(x)) * 0.1 |
| return (cover + residual).clamp(0, 1) |
|
|
|
|
|
|
| class StegDecoder(nn.Module): |
| def __init__(self, msg_bits=64, base=48): |
| super().__init__() |
| self.features = nn.Sequential( |
| nn.Conv2d(3, base, 3, padding=1), nn.BatchNorm2d(base), nn.ReLU(inplace=True), |
| nn.Conv2d(base, base * 2, 3, stride=2, padding=1), nn.BatchNorm2d(base * 2), nn.ReLU(inplace=True), |
| nn.Conv2d(base * 2, base * 4, 3, stride=2, padding=1), nn.BatchNorm2d(base * 4), nn.ReLU(inplace=True), |
| nn.Conv2d(base * 4, base * 4, 3, stride=2, padding=1), nn.BatchNorm2d(base * 4), nn.ReLU(inplace=True), |
| nn.AdaptiveAvgPool2d(1), |
| ) |
| self.head = nn.Linear(base * 4, msg_bits) |
|
|
| def forward(self, x): |
| f = self.features(x).flatten(1) |
| return self.head(f) |
|
|
