| import torch |
| import torch.nn as nn |
|
|
| class BaroNet(nn.Module): |
| def __init__(self): |
| super().__init__() |
|
|
| def conv_block(in_c, out_c): |
| return nn.Sequential( |
| nn.Conv2d(in_c, out_c, 3, padding=1), |
| nn.BatchNorm2d(out_c), |
| nn.ReLU(inplace=True), |
| nn.Conv2d(out_c, out_c, 3, padding=1), |
| nn.BatchNorm2d(out_c), |
| nn.ReLU(inplace=True) |
| ) |
|
|
| self.enc1 = conv_block(3, 64) |
| self.enc2 = conv_block(64, 128) |
| self.enc3 = conv_block(128, 256) |
| self.enc4 = conv_block(256, 512) |
|
|
| self.pool = nn.MaxPool2d(2) |
|
|
| self.up1 = nn.ConvTranspose2d(512, 256, 2, 2) |
| self.dec1 = conv_block(512, 256) |
|
|
| self.up2 = nn.ConvTranspose2d(256, 128, 2, 2) |
| self.dec2 = conv_block(256, 128) |
|
|
| self.up3 = nn.ConvTranspose2d(128, 64, 2, 2) |
| self.dec3 = conv_block(128, 64) |
|
|
| self.out = nn.Conv2d(64, 1, 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)) |
|
|
| d1 = self.up1(e4) |
| d1 = self.dec1(torch.cat([d1, e3], dim=1)) |
|
|
| d2 = self.up2(d1) |
| d2 = self.dec2(torch.cat([d2, e2], dim=1)) |
|
|
| d3 = self.up3(d2) |
| d3 = self.dec3(torch.cat([d3, e1], dim=1)) |
|
|
| return torch.sigmoid(self.out(d3)) |
