# from https://github.com/milesial/Pytorch-UNet/tree/master # This is a UNet model for CARRADA and KuRALS import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): """(convolution => [BN] => ReLU) * 2""" def __init__(self, in_channels, out_channels, mid_channels=None): super().__init__() if not mid_channels: mid_channels = out_channels self.double_conv = nn.Sequential( nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(mid_channels), nn.ReLU(inplace=True), nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True) ) def forward(self, x): return self.double_conv(x) class Down(nn.Module): """Downscaling with maxpool then double conv""" def __init__(self, in_channels, out_channels): super().__init__() self.maxpool_conv = nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_channels, out_channels) ) def forward(self, x): return self.maxpool_conv(x) class Up(nn.Module): """Upscaling then double conv""" def __init__(self, in_channels, out_channels, bilinear=True): super().__init__() # if bilinear, use the normal convolutions to reduce the number of channels if bilinear: self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) self.conv = DoubleConv(in_channels, out_channels, in_channels // 2) else: self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2) self.conv = DoubleConv(in_channels, out_channels) def forward(self, x1, x2): x1 = self.up(x1) # input is BCHW diffY = x2.size()[2] - x1.size()[2] diffX = x2.size()[3] - x1.size()[3] x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2]) # if you have padding issues, see # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd x = torch.cat([x2, x1], dim=1) return self.conv(x) class OutConv(nn.Module): def __init__(self, in_channels, out_channels): super(OutConv, self).__init__() self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, n_classes, n_frames=5, bilinear=True): """ Args: n_classes : _description_ n_frames : _description_. Defaults to 5. bilinear : Bilinear interpolation or ConvTranspose2d. """ super(UNet, self).__init__() self.n_channels = n_frames self.n_classes = n_classes self.bilinear = bilinear # lt@20240406 reduce resolution # self.pool0 = nn.MaxPool2d(2, stride=2, ceil_mode=True) self.inc = (DoubleConv(n_frames, 64)) self.down1 = (Down(64, 128)) self.down2 = (Down(128, 256)) self.down3 = (Down(256, 512)) factor = 2 if bilinear else 1 self.down4 = (Down(512, 1024 // factor)) self.up1 = (Up(1024, 512 // factor, bilinear)) self.up2 = (Up(512, 256 // factor, bilinear)) self.up3 = (Up(256, 128 // factor, bilinear)) self.up4 = (Up(128, 64, bilinear)) self.outc = (OutConv(64, n_classes)) def forward(self, x): # lt@20240406 reduce resolution # x = self.pool0(x) # lt@20240408 add resolution # x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True) x1 = self.inc(x) x2 = self.down1(x1) x3 = self.down2(x2) x4 = self.down3(x3) x5 = self.down4(x4) x = self.up1(x5, x4) x = self.up2(x, x3) x = self.up3(x, x2) x = self.up4(x, x1) logits = self.outc(x) # lt@20240406 recover resolution # logits = F.interpolate(logits, scale_factor=2, mode='bilinear', align_corners=True) # lt@20240408 add resolution # logits = F.interpolate(logits, scale_factor=0.5, mode='bilinear', align_corners=True) return logits # from https://github.com/milesial/Pytorch-UNet/tree/master # import torch # import torch.nn as nn # import torch.nn.functional as F # lt@20240315 spatial group normalization # lt@20240316 change fixed value 200 to w//4. # class DoubleConv(nn.Module): # """(convolution => [SGN] => ReLU) * 2""" # def __init__(self, in_channels, out_channels, mid_channels=None): # super().__init__() # if not mid_channels: # mid_channels = out_channels # self.conv1 = nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False) # self.bn1_1 = nn.BatchNorm2d(mid_channels) # self.bn1_2 = nn.BatchNorm2d(mid_channels) # self.relu = nn.ReLU(inplace=True) # self.conv2 = nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False) # self.bn2_1 = nn.BatchNorm2d(out_channels) # self.bn2_2 = nn.BatchNorm2d(out_channels) # def forward(self, x): # b, c, h, w = x.shape # x = self.conv1(x) # x11 = self.bn1_1(x[:, :, :, :(w//4)]) # x12 = self.bn1_2(x[:, :, :, (w//4):]) # x1 = self.relu(torch.concat([x11, x12], dim=-1)) # x2 = self.conv2(x1) # x21 = self.bn2_1(x2[:, :, :, :(w//4)]) # x22 = self.bn2_2(x2[:, :, :, (w//4):]) # x2 = self.relu(torch.concat([x21, x22], dim=-1)) # return x2 # lt@20240315 Group Normalization # class DoubleConv(nn.Module): # """(convolution => [GN] => ReLU) * 2""" # def __init__(self, in_channels, out_channels, mid_channels=None): # super().__init__() # if not mid_channels: # mid_channels = out_channels # self.double_conv = nn.Sequential( # nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False), # nn.GroupNorm(mid_channels//16, mid_channels), # nn.ReLU(inplace=True), # nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False), # nn.GroupNorm(out_channels//16, out_channels), # nn.ReLU(inplace=True) # ) # def forward(self, x): # return self.double_conv(x) # class Down(nn.Module): # """Downscaling with maxpool then double conv""" # def __init__(self, in_channels, out_channels): # super().__init__() # self.maxpool_conv = nn.Sequential( # nn.MaxPool2d(2), # DoubleConv(in_channels, out_channels) # ) # def forward(self, x): # return self.maxpool_conv(x) # class Up(nn.Module): # """Upscaling then double conv""" # def __init__(self, in_channels, out_channels, bilinear=True): # super().__init__() # # if bilinear, use the normal convolutions to reduce the number of channels # if bilinear: # self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) # self.conv = DoubleConv(in_channels, out_channels, in_channels // 2) # else: # self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2) # self.conv = DoubleConv(in_channels, out_channels) # def forward(self, x1, x2): # x1 = self.up(x1) # # input is BCHW # diffY = x2.size()[2] - x1.size()[2] # diffX = x2.size()[3] - x1.size()[3] # x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, # diffY // 2, diffY - diffY // 2]) # # if you have padding issues, see # # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a # # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd # x = torch.cat([x2, x1], dim=1) # return self.conv(x) # class OutConv(nn.Module): # def __init__(self, in_channels, out_channels): # super(OutConv, self).__init__() # self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1) # def forward(self, x): # return self.conv(x) # class UNet(nn.Module): # def __init__(self, n_classes, n_frames=5, bilinear=True): # """ # Args: # n_classes : _description_ # n_frames : _description_. Defaults to 5. # bilinear : Bilinear interpolation or ConvTranspose2d. # """ # super(UNet, self).__init__() # self.n_channels = n_frames # self.n_classes = n_classes # self.bilinear = bilinear # self.inc = (DoubleConv(n_frames, 64)) # self.down1 = (Down(64, 128)) # self.down2 = (Down(128, 256)) # self.down3 = (Down(256, 512)) # factor = 2 if bilinear else 1 # self.down4 = (Down(512, 1024 // factor)) # self.up1 = (Up(1024, 512 // factor, bilinear)) # self.up2 = (Up(512, 256 // factor, bilinear)) # self.up3 = (Up(256, 128 // factor, bilinear)) # self.up4 = (Up(128, 64, bilinear)) # self.outc = (OutConv(64, n_classes)) # def forward(self, x): # x1 = self.inc(x) # x2 = self.down1(x1) # x3 = self.down2(x2) # x4 = self.down3(x3) # x5 = self.down4(x4) # x = self.up1(x5, x4) # x = self.up2(x, x3) # x = self.up3(x, x2) # x = self.up4(x, x1) # logits = self.outc(x) # return logits