| import torch |
| import torch.nn as nn |
|
|
| from .vit import ( |
| _make_pretrained_vitb_rn50_384, |
| _make_pretrained_vitl16_384, |
| _make_pretrained_vitb16_384, |
| forward_vit, |
| ) |
|
|
|
|
| def _make_encoder( |
| backbone, |
| features, |
| use_pretrained, |
| groups=1, |
| expand=False, |
| exportable=True, |
| hooks=None, |
| use_vit_only=False, |
| use_readout="ignore", |
| enable_attention_hooks=False, |
| ): |
| if backbone == "vitl16_384": |
| pretrained = _make_pretrained_vitl16_384( |
| use_pretrained, |
| hooks=hooks, |
| use_readout=use_readout, |
| enable_attention_hooks=enable_attention_hooks, |
| ) |
| scratch = _make_scratch( |
| [256, 512, 1024, 1024], features, groups=groups, expand=expand |
| ) |
| elif backbone == "vitb_rn50_384": |
| pretrained = _make_pretrained_vitb_rn50_384( |
| use_pretrained, |
| hooks=hooks, |
| use_vit_only=use_vit_only, |
| use_readout=use_readout, |
| enable_attention_hooks=enable_attention_hooks, |
| ) |
| scratch = _make_scratch( |
| [256, 512, 768, 768], features, groups=groups, expand=expand |
| ) |
| elif backbone == "vitb16_384": |
| pretrained = _make_pretrained_vitb16_384( |
| use_pretrained, |
| hooks=hooks, |
| use_readout=use_readout, |
| enable_attention_hooks=enable_attention_hooks, |
| ) |
| scratch = _make_scratch( |
| [96, 192, 384, 768], features, groups=groups, expand=expand |
| ) |
| elif backbone == "resnext101_wsl": |
| pretrained = _make_pretrained_resnext101_wsl(use_pretrained) |
| scratch = _make_scratch( |
| [256, 512, 1024, 2048], features, groups=groups, expand=expand |
| ) |
| else: |
| print(f"Backbone '{backbone}' not implemented") |
| assert False |
|
|
| return pretrained, scratch |
|
|
|
|
| def _make_scratch(in_shape, out_shape, groups=1, expand=False): |
| scratch = nn.Module() |
|
|
| out_shape1 = out_shape |
| out_shape2 = out_shape |
| out_shape3 = out_shape |
| out_shape4 = out_shape |
| if expand == True: |
| out_shape1 = out_shape |
| out_shape2 = out_shape * 2 |
| out_shape3 = out_shape * 4 |
| out_shape4 = out_shape * 8 |
|
|
| scratch.layer1_rn = nn.Conv2d( |
| in_shape[0], |
| out_shape1, |
| kernel_size=3, |
| stride=1, |
| padding=1, |
| bias=False, |
| groups=groups, |
| ) |
| scratch.layer2_rn = nn.Conv2d( |
| in_shape[1], |
| out_shape2, |
| kernel_size=3, |
| stride=1, |
| padding=1, |
| bias=False, |
| groups=groups, |
| ) |
| scratch.layer3_rn = nn.Conv2d( |
| in_shape[2], |
| out_shape3, |
| kernel_size=3, |
| stride=1, |
| padding=1, |
| bias=False, |
| groups=groups, |
| ) |
| scratch.layer4_rn = nn.Conv2d( |
| in_shape[3], |
| out_shape4, |
| kernel_size=3, |
| stride=1, |
| padding=1, |
| bias=False, |
| groups=groups, |
| ) |
|
|
| return scratch |
|
|
|
|
| def _make_resnet_backbone(resnet): |
| pretrained = nn.Module() |
| pretrained.layer1 = nn.Sequential( |
| resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool, resnet.layer1 |
| ) |
|
|
| pretrained.layer2 = resnet.layer2 |
| pretrained.layer3 = resnet.layer3 |
| pretrained.layer4 = resnet.layer4 |
|
|
| return pretrained |
|
|
|
|
| def _make_pretrained_resnext101_wsl(use_pretrained): |
| resnet = torch.hub.load("facebookresearch/WSL-Images", "resnext101_32x8d_wsl") |
| return _make_resnet_backbone(resnet) |
|
|
|
|
| class Interpolate(nn.Module): |
| """Interpolation module.""" |
|
|
| def __init__(self, scale_factor, mode, align_corners=False): |
| """Init. |
| |
| Args: |
| scale_factor (float): scaling |
| mode (str): interpolation mode |
| """ |
| super(Interpolate, self).__init__() |
|
|
| self.interp = nn.functional.interpolate |
| self.scale_factor = scale_factor |
| self.mode = mode |
| self.align_corners = align_corners |
|
|
| def forward(self, x): |
| """Forward pass. |
| |
| Args: |
| x (tensor): input |
| |
| Returns: |
| tensor: interpolated data |
| """ |
|
|
| x = self.interp( |
| x, |
| scale_factor=self.scale_factor, |
| mode=self.mode, |
| align_corners=self.align_corners, |
| ) |
|
|
| return x |
|
|
|
|
| class ResidualConvUnit(nn.Module): |
| """Residual convolution module.""" |
|
|
| def __init__(self, features): |
| """Init. |
| |
| Args: |
| features (int): number of features |
| """ |
| super().__init__() |
|
|
| self.conv1 = nn.Conv2d( |
| features, features, kernel_size=3, stride=1, padding=1, bias=True |
| ) |
|
|
| self.conv2 = nn.Conv2d( |
| features, features, kernel_size=3, stride=1, padding=1, bias=True |
| ) |
|
|
| self.relu = nn.ReLU(inplace=True) |
|
|
| def forward(self, x): |
| """Forward pass. |
| |
| Args: |
| x (tensor): input |
| |
| Returns: |
| tensor: output |
| """ |
| out = self.relu(x) |
| out = self.conv1(out) |
| out = self.relu(out) |
| out = self.conv2(out) |
|
|
| return out + x |
|
|
|
|
| class FeatureFusionBlock(nn.Module): |
| """Feature fusion block.""" |
|
|
| def __init__(self, features): |
| """Init. |
| |
| Args: |
| features (int): number of features |
| """ |
| super(FeatureFusionBlock, self).__init__() |
|
|
| self.resConfUnit1 = ResidualConvUnit(features) |
| self.resConfUnit2 = ResidualConvUnit(features) |
|
|
| def forward(self, *xs): |
| """Forward pass. |
| |
| Returns: |
| tensor: output |
| """ |
| output = xs[0] |
|
|
| if len(xs) == 2: |
| output += self.resConfUnit1(xs[1]) |
|
|
| output = self.resConfUnit2(output) |
|
|
| output = nn.functional.interpolate( |
| output, scale_factor=2, mode="bilinear", align_corners=True |
| ) |
|
|
| return output |
|
|
|
|
| class ResidualConvUnit_custom(nn.Module): |
| """Residual convolution module.""" |
|
|
| def __init__(self, features, activation, bn): |
| """Init. |
| |
| Args: |
| features (int): number of features |
| """ |
| super().__init__() |
|
|
| self.bn = bn |
|
|
| self.groups = 1 |
|
|
| self.conv1 = nn.Conv2d( |
| features, |
| features, |
| kernel_size=3, |
| stride=1, |
| padding=1, |
| bias=not self.bn, |
| groups=self.groups, |
| ) |
|
|
| self.conv2 = nn.Conv2d( |
| features, |
| features, |
| kernel_size=3, |
| stride=1, |
| padding=1, |
| bias=not self.bn, |
| groups=self.groups, |
| ) |
|
|
| if self.bn == True: |
| self.bn1 = nn.BatchNorm2d(features) |
| self.bn2 = nn.BatchNorm2d(features) |
|
|
| self.activation = activation |
|
|
| self.skip_add = nn.quantized.FloatFunctional() |
|
|
| def forward(self, x): |
| """Forward pass. |
| |
| Args: |
| x (tensor): input |
| |
| Returns: |
| tensor: output |
| """ |
|
|
| out = self.activation(x) |
| out = self.conv1(out) |
| if self.bn == True: |
| out = self.bn1(out) |
|
|
| out = self.activation(out) |
| out = self.conv2(out) |
| if self.bn == True: |
| out = self.bn2(out) |
|
|
| if self.groups > 1: |
| out = self.conv_merge(out) |
|
|
| return self.skip_add.add(out, x) |
|
|
| |
|
|
|
|
| class FeatureFusionBlock_custom(nn.Module): |
| """Feature fusion block.""" |
|
|
| def __init__( |
| self, |
| features, |
| activation, |
| deconv=False, |
| bn=False, |
| expand=False, |
| align_corners=True, |
| ): |
| """Init. |
| |
| Args: |
| features (int): number of features |
| """ |
| super(FeatureFusionBlock_custom, self).__init__() |
|
|
| self.deconv = deconv |
| self.align_corners = align_corners |
|
|
| self.groups = 1 |
|
|
| self.expand = expand |
| out_features = features |
| if self.expand == True: |
| out_features = features // 2 |
|
|
| self.out_conv = nn.Conv2d( |
| features, |
| out_features, |
| kernel_size=1, |
| stride=1, |
| padding=0, |
| bias=True, |
| groups=1, |
| ) |
|
|
| self.resConfUnit1 = ResidualConvUnit_custom(features, activation, bn) |
| self.resConfUnit2 = ResidualConvUnit_custom(features, activation, bn) |
|
|
| self.skip_add = nn.quantized.FloatFunctional() |
|
|
| def forward(self, *xs): |
| """Forward pass. |
| |
| Returns: |
| tensor: output |
| """ |
| output = xs[0] |
|
|
| if len(xs) == 2: |
| res = self.resConfUnit1(xs[1]) |
| output = self.skip_add.add(output, res) |
| |
|
|
| output = self.resConfUnit2(output) |
|
|
| output = nn.functional.interpolate( |
| output, scale_factor=2, mode="bilinear", align_corners=self.align_corners |
| ) |
|
|
| output = self.out_conv(output) |
|
|
| return output |
|
|
