|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
|
|
|
class ResidualBlock(nn.Module): |
|
|
def __init__(self, in_channels, out_channels): |
|
|
super(ResidualBlock, self).__init__() |
|
|
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False) |
|
|
self.bn1 = nn.BatchNorm2d(out_channels) |
|
|
self.relu = nn.ReLU(inplace=True) |
|
|
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False) |
|
|
self.bn2 = nn.BatchNorm2d(out_channels) |
|
|
if in_channels != out_channels: |
|
|
self.shortcut = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False), nn.BatchNorm2d(out_channels)) |
|
|
else: |
|
|
self.shortcut = nn.Identity() |
|
|
def forward(self, x): |
|
|
residual = self.shortcut(x) |
|
|
out = self.conv1(x); out = self.bn1(out); out = self.relu(out) |
|
|
out = self.conv2(out); out = self.bn2(out) |
|
|
out += residual |
|
|
out = self.relu(out) |
|
|
return out |
|
|
|
|
|
class ComplexUNet(nn.Module): |
|
|
def __init__(self, base_channels=96): |
|
|
super(ComplexUNet, self).__init__() |
|
|
c = base_channels |
|
|
self.pool = nn.MaxPool2d(2, 2) |
|
|
self.enc1 = ResidualBlock(3, c) |
|
|
self.enc2 = ResidualBlock(c, c*2) |
|
|
self.enc3 = ResidualBlock(c*2, c*4) |
|
|
self.enc4 = ResidualBlock(c*4, c*8) |
|
|
self.bottleneck = ResidualBlock(c*8, c*16) |
|
|
self.upconv1 = nn.ConvTranspose2d(c*16, c*8, kernel_size=2, stride=2) |
|
|
self.upconv2 = nn.ConvTranspose2d(c*8, c*4, kernel_size=2, stride=2) |
|
|
self.upconv3 = nn.ConvTranspose2d(c*4, c*2, kernel_size=2, stride=2) |
|
|
self.upconv4 = nn.ConvTranspose2d(c*2, c, kernel_size=2, stride=2) |
|
|
self.dec_conv1 = ResidualBlock(c*16, c*8) |
|
|
self.dec_conv2 = ResidualBlock(c*8, c*4) |
|
|
self.dec_conv3 = ResidualBlock(c*4, c*2) |
|
|
self.dec_conv4 = ResidualBlock(c*2, c) |
|
|
self.final_conv = nn.Conv2d(c, 3, kernel_size=1) |
|
|
def forward(self, x): |
|
|
e1 = self.enc1(x); p1 = self.pool(e1); e2 = self.enc2(p1); p2 = self.pool(e2) |
|
|
e3 = self.enc3(p2); p3 = self.pool(e3); e4 = self.enc4(p3); p4 = self.pool(e4) |
|
|
b = self.bottleneck(p4) |
|
|
d1 = self.upconv1(b); d1 = torch.cat([d1, e4], dim=1); d1 = self.dec_conv1(d1) |
|
|
d2 = self.upconv2(d1); d2 = torch.cat([d2, e3], dim=1); d2 = self.dec_conv2(d2) |
|
|
d3 = self.upconv3(d2); d3 = torch.cat([d3, e2], dim=1); d3 = self.dec_conv3(d3) |
|
|
d4 = self.upconv4(d3); d4 = torch.cat([d4, e1], dim=1); d4 = self.dec_conv4(d4) |
|
|
out = self.final_conv(d4) |
|
|
return torch.sigmoid(out) |
|
|
|
