Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import torchvision.models as models | |
| class UpsamplingAdd(nn.Module): | |
| def __init__(self, in_channels: int, out_channels: int, scale_factor: int = 2): | |
| super().__init__() | |
| self.upsample_layer = nn.Sequential( | |
| nn.Upsample( | |
| scale_factor=scale_factor, mode="bilinear", align_corners=False | |
| ), | |
| nn.Conv2d(in_channels, out_channels, | |
| kernel_size=1, padding=0, bias=False), | |
| nn.InstanceNorm2d(out_channels), | |
| ) | |
| def forward(self, x: torch.Tensor, x_skip: torch.Tensor): | |
| # Check if the width dimension is odd and needs zero padding | |
| x = self.upsample_layer(x) | |
| if x.shape[-1] != x_skip.shape[-1] or x.shape[-2] != x_skip.shape[-2]: | |
| x = nn.functional.interpolate( | |
| x, size=(x_skip.shape[-2], x_skip.shape[-1]), mode="bilinear" | |
| ) | |
| return x + x_skip | |
| class SegmentationHead(nn.Module): | |
| def __init__(self, in_channels: int, n_classes: int, dropout_rate: float = 0.0): | |
| super(SegmentationHead, self).__init__() | |
| backbone = models.resnet18(pretrained=False, zero_init_residual=True) | |
| self.first_conv = nn.Conv2d( | |
| in_channels, 64, kernel_size=7, stride=2, padding=3, bias=False | |
| ) | |
| self.bn1 = backbone.bn1 | |
| self.relu = backbone.relu | |
| self.layer1 = backbone.layer1 | |
| self.layer2 = backbone.layer2 | |
| self.layer3 = backbone.layer3 | |
| # Upsampling layers | |
| self.up3_skip = UpsamplingAdd( | |
| in_channels=256, out_channels=128, scale_factor=2) | |
| self.up2_skip = UpsamplingAdd( | |
| in_channels=128, out_channels=64, scale_factor=2) | |
| self.up1_skip = UpsamplingAdd( | |
| in_channels=64, out_channels=in_channels, scale_factor=2) | |
| # Segmentation head | |
| self.dropout = nn.Dropout( | |
| dropout_rate) if dropout_rate > 0 else nn.Identity() | |
| self.segmentation_head = nn.Sequential( | |
| nn.Conv2d(in_channels, in_channels, | |
| kernel_size=3, padding=1, bias=False), | |
| nn.InstanceNorm2d(in_channels), | |
| nn.ReLU(inplace=True), | |
| self.dropout, | |
| nn.Conv2d(in_channels, n_classes, kernel_size=1, padding=0), | |
| ) | |
| def forward(self, x: torch.Tensor): | |
| # (H, W) | |
| skip_x = {"1": x} | |
| x = self.first_conv(x) | |
| x = self.bn1(x) | |
| x = self.relu(x) | |
| x = self.dropout(x) | |
| # (H/4, W/4) | |
| x = self.layer1(x) | |
| skip_x["2"] = x | |
| x = self.dropout(x) | |
| x = self.layer2(x) | |
| skip_x["3"] = x | |
| x = self.dropout(x) | |
| # (H/8, W/8) | |
| x = self.layer3(x) | |
| x = self.dropout(x) | |
| # First upsample to (H/4, W/4) | |
| x = self.up3_skip(x, skip_x["3"]) | |
| x = self.dropout(x) | |
| # Second upsample to (H/2, W/2) | |
| x = self.up2_skip(x, skip_x["2"]) | |
| x = self.dropout(x) | |
| # Third upsample to (H, W) | |
| x = self.up1_skip(x, skip_x["1"]) | |
| x = self.dropout(x) | |
| segmentation_output = self.segmentation_head(x) | |
| return segmentation_output | |