Spaces:
Sleeping
Sleeping
| """ | |
| DenseNet encoder from ICAL. | |
| Copied from: https://github.com/qingzhenduyu/ICAL | |
| """ | |
| import math | |
| from typing import Tuple | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from einops import rearrange | |
| from torch import FloatTensor, LongTensor | |
| from .pos_enc import ImgPosEnc | |
| # DenseNet-B | |
| class _Bottleneck(nn.Module): | |
| def __init__(self, n_channels: int, growth_rate: int, use_dropout: bool): | |
| super(_Bottleneck, self).__init__() | |
| interChannels = 4 * growth_rate | |
| self.bn1 = nn.BatchNorm2d(interChannels) | |
| self.conv1 = nn.Conv2d(n_channels, interChannels, | |
| kernel_size=1, bias=False) | |
| self.bn2 = nn.BatchNorm2d(growth_rate) | |
| self.conv2 = nn.Conv2d( | |
| interChannels, growth_rate, kernel_size=3, padding=1, bias=False | |
| ) | |
| self.use_dropout = use_dropout | |
| self.dropout = nn.Dropout(p=0.2) | |
| def forward(self, x): | |
| out = F.relu(self.bn1(self.conv1(x)), inplace=True) | |
| if self.use_dropout: | |
| out = self.dropout(out) | |
| out = F.relu(self.bn2(self.conv2(out)), inplace=True) | |
| if self.use_dropout: | |
| out = self.dropout(out) | |
| out = torch.cat((x, out), 1) | |
| return out | |
| # single layer | |
| class _SingleLayer(nn.Module): | |
| def __init__(self, n_channels: int, growth_rate: int, use_dropout: bool): | |
| super(_SingleLayer, self).__init__() | |
| self.bn1 = nn.BatchNorm2d(n_channels) | |
| self.conv1 = nn.Conv2d( | |
| n_channels, growth_rate, kernel_size=3, padding=1, bias=False | |
| ) | |
| self.use_dropout = use_dropout | |
| self.dropout = nn.Dropout(p=0.2) | |
| def forward(self, x): | |
| out = self.conv1(F.relu(x, inplace=True)) | |
| if self.use_dropout: | |
| out = self.dropout(out) | |
| out = torch.cat((x, out), 1) | |
| return out | |
| # transition layer | |
| class _Transition(nn.Module): | |
| def __init__(self, n_channels: int, n_out_channels: int, use_dropout: bool): | |
| super(_Transition, self).__init__() | |
| self.bn1 = nn.BatchNorm2d(n_out_channels) | |
| self.conv1 = nn.Conv2d(n_channels, n_out_channels, | |
| kernel_size=1, bias=False) | |
| self.use_dropout = use_dropout | |
| self.dropout = nn.Dropout(p=0.2) | |
| def forward(self, x): | |
| out = F.relu(self.bn1(self.conv1(x)), inplace=True) | |
| if self.use_dropout: | |
| out = self.dropout(out) | |
| out = F.avg_pool2d(out, 2, ceil_mode=True) | |
| return out | |
| class DenseNet(nn.Module): | |
| def __init__( | |
| self, | |
| growth_rate: int, | |
| num_layers: int, | |
| reduction: float = 0.5, | |
| bottleneck: bool = True, | |
| use_dropout: bool = True, | |
| ): | |
| super(DenseNet, self).__init__() | |
| n_dense_blocks = num_layers | |
| n_channels = 2 * growth_rate | |
| self.conv1 = nn.Conv2d( | |
| 1, n_channels, kernel_size=7, padding=3, stride=2, bias=False | |
| ) | |
| self.norm1 = nn.BatchNorm2d(n_channels) | |
| self.dense1 = self._make_dense( | |
| n_channels, growth_rate, n_dense_blocks, bottleneck, use_dropout | |
| ) | |
| n_channels += n_dense_blocks * growth_rate | |
| n_out_channels = int(math.floor(n_channels * reduction)) | |
| self.trans1 = _Transition(n_channels, n_out_channels, use_dropout) | |
| n_channels = n_out_channels | |
| self.dense2 = self._make_dense( | |
| n_channels, growth_rate, n_dense_blocks, bottleneck, use_dropout | |
| ) | |
| n_channels += n_dense_blocks * growth_rate | |
| n_out_channels = int(math.floor(n_channels * reduction)) | |
| self.trans2 = _Transition(n_channels, n_out_channels, use_dropout) | |
| n_channels = n_out_channels | |
| self.dense3 = self._make_dense( | |
| n_channels, growth_rate, n_dense_blocks, bottleneck, use_dropout | |
| ) | |
| self.out_channels = n_channels + n_dense_blocks * growth_rate | |
| self.post_norm = nn.BatchNorm2d(self.out_channels) | |
| def _make_dense(n_channels, growth_rate, n_dense_blocks, bottleneck, use_dropout): | |
| layers = [] | |
| for _ in range(int(n_dense_blocks)): | |
| if bottleneck: | |
| layers.append(_Bottleneck( | |
| n_channels, growth_rate, use_dropout)) | |
| else: | |
| layers.append(_SingleLayer( | |
| n_channels, growth_rate, use_dropout)) | |
| n_channels += growth_rate | |
| return nn.Sequential(*layers) | |
| def forward(self, x, x_mask): | |
| out = self.conv1(x) | |
| out = self.norm1(out) | |
| out_mask = x_mask[:, 0::2, 0::2] | |
| out = F.relu(out, inplace=True) | |
| out = F.max_pool2d(out, 2, ceil_mode=True) | |
| out_mask = out_mask[:, 0::2, 0::2] | |
| out = self.dense1(out) | |
| out = self.trans1(out) | |
| out_mask = out_mask[:, 0::2, 0::2] | |
| out = self.dense2(out) | |
| out = self.trans2(out) | |
| out_mask = out_mask[:, 0::2, 0::2] | |
| out = self.dense3(out) | |
| out = self.post_norm(out) | |
| return out, out_mask | |
| class Encoder(nn.Module): | |
| def __init__(self, d_model: int, growth_rate: int, num_layers: int): | |
| super().__init__() | |
| self.model = DenseNet(growth_rate=growth_rate, num_layers=num_layers) | |
| self.feature_proj = nn.Conv2d( | |
| self.model.out_channels, d_model, kernel_size=1) | |
| self.pos_enc_2d = ImgPosEnc(d_model, normalize=True) | |
| self.norm = nn.LayerNorm(d_model) | |
| def forward( | |
| self, img: FloatTensor, img_mask: LongTensor | |
| ) -> Tuple[FloatTensor, LongTensor]: | |
| """encode image to feature | |
| Parameters | |
| ---------- | |
| img : FloatTensor | |
| [b, 1, h', w'] | |
| img_mask: LongTensor | |
| [b, h', w'] | |
| Returns | |
| ------- | |
| Tuple[FloatTensor, LongTensor] | |
| [b, h, w, d], [b, h, w] | |
| """ | |
| # extract feature | |
| feature, mask = self.model(img, img_mask) | |
| feature = self.feature_proj(feature) | |
| # proj | |
| feature = rearrange(feature, "b d h w -> b h w d") | |
| # positional encoding | |
| feature = self.pos_enc_2d(feature, mask) | |
| feature = self.norm(feature) | |
| return feature, mask | |