| class DoubleConv(nn.Module): | |
| def __init__(self, in_ch, out_ch): | |
| super().__init__() | |
| self.block = nn.Sequential( | |
| nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False), | |
| nn.BatchNorm2d(out_ch), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), | |
| nn.BatchNorm2d(out_ch), | |
| nn.ReLU(inplace=True), | |
| ) | |
| def forward(self, x): | |
| return self.block(x) | |
| class UNet(nn.Module): | |
| def __init__(self, in_ch=3, out_ch=1, base=64): | |
| super().__init__() | |
| c = [base, base * 2, base * 4, base * 8, base * 16] | |
| self.enc1 = DoubleConv(in_ch, c[0]) | |
| self.enc2 = DoubleConv(c[0], c[1]) | |
| self.enc3 = DoubleConv(c[1], c[2]) | |
| self.enc4 = DoubleConv(c[2], c[3]) | |
| self.bottleneck = DoubleConv(c[3], c[4]) | |
| self.pool = nn.MaxPool2d(2) | |
| self.up4 = nn.ConvTranspose2d(c[4], c[3], 2, stride=2) | |
| self.dec4 = DoubleConv(c[4], c[3]) | |
| self.up3 = nn.ConvTranspose2d(c[3], c[2], 2, stride=2) | |
| self.dec3 = DoubleConv(c[3], c[2]) | |
| self.up2 = nn.ConvTranspose2d(c[2], c[1], 2, stride=2) | |
| self.dec2 = DoubleConv(c[2], c[1]) | |
| self.up1 = nn.ConvTranspose2d(c[1], c[0], 2, stride=2) | |
| self.dec1 = DoubleConv(c[1], c[0]) | |
| self.head = nn.Conv2d(c[0], out_ch, 1) | |
| def forward(self, x): | |
| e1 = self.enc1(x) | |
| e2 = self.enc2(self.pool(e1)) | |
| e3 = self.enc3(self.pool(e2)) | |
| e4 = self.enc4(self.pool(e3)) | |
| b = self.bottleneck(self.pool(e4)) | |
| d4 = self.dec4(torch.cat([self.up4(b), e4], 1)) | |
| d3 = self.dec3(torch.cat([self.up3(d4), e3], 1)) | |
| d2 = self.dec2(torch.cat([self.up2(d3), e2], 1)) | |
| d1 = self.dec1(torch.cat([self.up1(d2), e1], 1)) | |
| return self.head(d1) |