| |
| import warnings |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from mmcv.cnn import build_upsample_layer |
|
|
| BN_MOMENTUM = 0.1 |
| BN3D_MOMENTUM = 0.05 |
|
|
|
|
| def point_wise_block(inplanes, outplanes): |
| return nn.Sequential( |
| nn.Conv2d( |
| in_channels=inplanes, |
| out_channels=outplanes, |
| kernel_size=1, |
| padding=0, |
| stride=1, |
| bias=False, |
| ), |
| nn.BatchNorm2d(outplanes, momentum=BN_MOMENTUM), |
| nn.ReLU(inplace=True), |
| ) |
|
|
|
|
| def conv_block(inplanes, outplanes, kernel_size, stride=1, padding=0): |
| return nn.Sequential( |
| nn.Conv2d( |
| in_channels=inplanes, |
| out_channels=outplanes, |
| kernel_size=kernel_size, |
| padding=padding, |
| stride=stride, |
| bias=False, |
| ), |
| nn.BatchNorm2d(outplanes, momentum=BN_MOMENTUM), |
| nn.ReLU(inplace=True), |
| ) |
|
|
|
|
| def conv3x3(in_planes, out_planes, stride=1): |
| """3x3 convolution with padding""" |
| return nn.Conv2d( |
| in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False |
| ) |
|
|
|
|
| def conv3d_block( |
| inplanes, |
| outplanes, |
| kernel_size=(3, 1, 1), |
| stride=(1, 1, 1), |
| padding=0, |
| bias=False, |
| inplace=False, |
| act=nn.GELU, |
| ): |
| """ |
| General conv3d block for handling 3d feature maps |
| """ |
| return nn.Sequential( |
| nn.Conv3d( |
| inplanes, |
| outplanes, |
| kernel_size=kernel_size, |
| stride=stride, |
| padding=padding, |
| bias=bias, |
| ), |
| nn.BatchNorm3d(outplanes, momentum=BN3D_MOMENTUM), |
| act(), |
| ) |
|
|
|
|
| def deconv3d_block( |
| inplanes, |
| outplanes, |
| kernel_size=(2, 4, 4), |
| stride=(2, 2, 2), |
| padding=(0, 1, 1), |
| bias=False, |
| inplace=False, |
| out_padding=(0, 1, 1), |
| act=None, |
| ): |
| """ |
| General Transpose 3D Convolution for handling 3D feature maps |
| """ |
| layers = [] |
| layers.append( |
| build_upsample_layer( |
| dict(type="deconv3d"), |
| in_channels=inplanes, |
| out_channels=outplanes, |
| kernel_size=kernel_size, |
| stride=stride, |
| padding=padding, |
| output_padding=out_padding, |
| bias=bias, |
| ) |
| ) |
| layers.append(nn.BatchNorm3d(outplanes)) |
|
|
| if act is not None: |
| layers.append(act()) |
|
|
| return nn.Sequential(*layers) |
|
|
|
|
| def resize( |
| input, |
| size=None, |
| scale_factor=None, |
| mode="nearest", |
| align_corners=None, |
| warning=True, |
| ): |
| if warning: |
| if size is not None and align_corners: |
| input_h, input_w = tuple(int(x) for x in input.shape[2:]) |
| output_h, output_w = tuple(int(x) for x in size) |
| if output_h > input_h or output_w > output_h: |
| if ( |
| (output_h > 1 and output_w > 1 and input_h > 1 and input_w > 1) |
| and (output_h - 1) % (input_h - 1) |
| and (output_w - 1) % (input_w - 1) |
| ): |
| warnings.warn( |
| f"When align_corners={align_corners}, " |
| "the output would more aligned if " |
| f"input size {(input_h, input_w)} is `x+1` and " |
| f"out size {(output_h, output_w)} is `nx+1`" |
| ) |
| if isinstance(size, torch.Size): |
| size = tuple(int(x) for x in size) |
| return F.interpolate(input, size, scale_factor, mode, align_corners) |
|
|
|
|
| class InceptionBlock(nn.Module): |
| def __init__(self, inplanes, outplanes, stride=1, pool_size=3): |
| self.inplanes = inplanes |
| self.outplanes = outplanes |
| self.stride = stride |
| self.pool_size = pool_size |
| super(InceptionBlock, self).__init__() |
|
|
| self.pw_block = point_wise_block(self.inplanes, self.outplanes // 4) |
| self.mp_layer = nn.MaxPool2d( |
| kernel_size=self.pool_size, stride=stride, padding=1 |
| ) |
| self.conv3_block = conv_block( |
| self.outplanes // 4, self.outplanes // 4, kernel_size=3, stride=1, padding=1 |
| ) |
| self.conv5_block = conv_block( |
| self.outplanes // 4, self.outplanes // 4, kernel_size=5, stride=1, padding=2 |
| ) |
|
|
| def forward(self, x): |
| x1 = self.pw_block(x) |
|
|
| x2 = self.pw_block(x) |
| x2 = self.conv3_block(x2) |
|
|
| x3 = self.pw_block(x) |
| x3 = self.conv5_block(x3) |
|
|
| x4 = self.mp_layer(x) |
| x4 = self.pw_block(x4) |
|
|
| x = torch.cat((x1, x2, x3, x4), dim=1) |
| return x |
|
|
|
|
| class InceptionBlock3D(nn.Module): |
| def __init__(self, inplanes, outplanes, stride=1, pool_size=3): |
| self.inplanes = inplanes |
| self.outplanes = outplanes |
| self.stride = stride |
| self.pool_size = pool_size |
| super(InceptionBlock3D, self).__init__() |
|
|
| self.pw_block = conv3d_block( |
| self.inplanes, |
| self.outplanes // 4, |
| kernel_size=(1, 1, 1), |
| stride=(1, 1, 1), |
| act=nn.ReLU, |
| ) |
| self.mp_layer = nn.MaxPool3d( |
| kernel_size=(self.pool_size, 1, 1), stride=stride, padding=(1, 0, 0) |
| ) |
| self.conv3_block = conv3d_block( |
| self.outplanes // 4, |
| self.outplanes // 4, |
| kernel_size=(3, 1, 1), |
| stride=1, |
| padding=(1, 0, 0), |
| act=nn.ReLU, |
| ) |
| self.conv5_block = conv3d_block( |
| self.outplanes // 4, |
| self.outplanes // 4, |
| kernel_size=(5, 1, 1), |
| stride=1, |
| padding=(2, 0, 0), |
| act=nn.ReLU, |
| ) |
|
|
| def forward(self, x): |
| x1 = self.pw_block(x) |
|
|
| x2 = self.pw_block(x) |
| x2 = self.conv3_block(x2) |
|
|
| x3 = self.pw_block(x) |
| x3 = self.conv5_block(x3) |
|
|
| x4 = self.mp_layer(x) |
| x4 = self.pw_block(x4) |
|
|
| x = torch.cat((x1, x2, x3, x4), dim=1) |
| return x |
|
|
|
|
| class SELayer(nn.Module): |
| def __init__(self, channel, reduction=16): |
| super(SELayer, self).__init__() |
| self.avg_pool = nn.AdaptiveAvgPool2d(1) |
| self.fc = nn.Sequential( |
| nn.Linear(channel, channel // reduction, bias=False), |
| nn.ReLU(inplace=True), |
| nn.Linear(channel // reduction, channel, bias=False), |
| nn.Sigmoid(), |
| ) |
|
|
| def forward(self, x): |
| b, c, _, _ = x.size() |
| y = self.avg_pool(x).view(b, c) |
| y = self.fc(y).view(b, c, 1, 1) |
| return x * y.expand_as(x) |
|
|
|
|
| class Texture_Enhance(nn.Module): |
| def __init__(self, num_features): |
| super().__init__() |
| |
| self.output_features = num_features * 4 |
| self.output_features_d = num_features |
| self.conv0 = nn.Conv2d(num_features, num_features, 1) |
| self.conv1 = nn.Conv2d(num_features, num_features, 3, padding=1) |
| self.bn1 = nn.BatchNorm2d(num_features) |
| self.conv2 = nn.Conv2d(num_features * 2, num_features, 3, padding=1) |
| self.bn2 = nn.BatchNorm2d(2 * num_features) |
| self.conv3 = nn.Conv2d(num_features * 3, num_features, 3, padding=1) |
| self.bn3 = nn.BatchNorm2d(3 * num_features) |
| self.conv_last = nn.Conv2d(num_features * 4, num_features * 4, 1) |
| self.bn4 = nn.BatchNorm2d(4 * num_features) |
| self.bn_last = nn.BatchNorm2d(num_features * 4) |
|
|
| def forward(self, feature_maps, attention_maps=(1, 1)): |
| B, N, H, W = feature_maps.shape |
|
|
| if type(attention_maps) == tuple: |
| attention_size = (int(H * attention_maps[0]), int(W * attention_maps[1])) |
| else: |
| attention_size = (attention_maps.shape[2], attention_maps.shape[3]) |
|
|
| feature_maps_d = F.adaptive_avg_pool2d(feature_maps, attention_size) |
| feature_maps = feature_maps - F.interpolate( |
| feature_maps_d, |
| (feature_maps.shape[2], feature_maps.shape[3]), |
| mode="nearest", |
| ) |
| feature_maps0 = self.conv0(feature_maps) |
| feature_maps1 = self.conv1(F.relu(self.bn1(feature_maps0), inplace=True)) |
| feature_maps1_ = torch.cat([feature_maps0, feature_maps1], dim=1) |
| feature_maps2 = self.conv2(F.relu(self.bn2(feature_maps1_), inplace=True)) |
| feature_maps2_ = torch.cat([feature_maps1_, feature_maps2], dim=1) |
| feature_maps3 = self.conv3(F.relu(self.bn3(feature_maps2_), inplace=True)) |
| feature_maps3_ = torch.cat([feature_maps2_, feature_maps3], dim=1) |
| feature_maps = self.bn_last( |
| self.conv_last(F.relu(self.bn4(feature_maps3_), inplace=True)) |
| ) |
| return feature_maps |
|
|
