File size: 1,559 Bytes
1fd7cfd |
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 |
import torch
import torch.nn as nn
class WatermarkRemover(nn.Module):
def __init__(self):
super(WatermarkRemover, self).__init__()
self.enc1 = self.conv_block(3, 64)
self.enc2 = self.conv_block(64, 128)
self.enc3 = self.conv_block(128, 256)
self.enc4 = self.conv_block(256, 512)
self.bottleneck = self.conv_block(512, 1024)
self.dec4 = self.conv_block(1024 + 512, 512)
self.dec3 = self.conv_block(512 + 256, 256)
self.dec2 = self.conv_block(256 + 128, 128)
self.dec1 = self.conv_block(128 + 64, 64)
self.final_layer = nn.Conv2d(64, 3, kernel_size=1)
def conv_block(self, in_channels, out_channels):
return nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
)
def forward(self, x):
e1 = self.enc1(x)
e2 = self.enc2(nn.MaxPool2d(2)(e1))
e3 = self.enc3(nn.MaxPool2d(2)(e2))
e4 = self.enc4(nn.MaxPool2d(2)(e3))
b = self.bottleneck(nn.MaxPool2d(2)(e4))
d4 = self.dec4(torch.cat((nn.Upsample(scale_factor=2)(b), e4), dim=1))
d3 = self.dec3(torch.cat((nn.Upsample(scale_factor=2)(d4), e3), dim=1))
d2 = self.dec2(torch.cat((nn.Upsample(scale_factor=2)(d3), e2), dim=1))
d1 = self.dec1(torch.cat((nn.Upsample(scale_factor=2)(d2), e1), dim=1))
return self.final_layer(d1)
|