Spaces:
Running
Running
| import numpy as np | |
| import torch.nn as nn | |
| import torch | |
| import torch.nn.functional as F | |
| class MLP(nn.Module): | |
| """ | |
| Linear Embedding: | |
| """ | |
| def __init__(self, input_dim=2048, embed_dim=768): | |
| super().__init__() | |
| self.proj = nn.Linear(input_dim, embed_dim) | |
| def forward(self, x): | |
| x = x.flatten(2).transpose(1, 2) | |
| x = self.proj(x) | |
| return x | |
| class DecoderHead(nn.Module): | |
| def __init__(self, | |
| in_channels=[64, 128, 320, 512], | |
| num_classes=40, | |
| dropout_ratio=0.1, | |
| norm_layer=nn.BatchNorm2d, | |
| embed_dim=768, | |
| align_corners=False): | |
| super(DecoderHead, self).__init__() | |
| self.num_classes = num_classes | |
| self.dropout_ratio = dropout_ratio | |
| self.align_corners = align_corners | |
| self.in_channels = in_channels | |
| if dropout_ratio > 0: | |
| self.dropout = nn.Dropout2d(dropout_ratio) | |
| else: | |
| self.dropout = None | |
| c1_in_channels, c2_in_channels, c3_in_channels, c4_in_channels = self.in_channels | |
| embedding_dim = embed_dim | |
| self.linear_c4 = MLP(input_dim=c4_in_channels, embed_dim=embedding_dim) | |
| self.linear_c3 = MLP(input_dim=c3_in_channels, embed_dim=embedding_dim) | |
| self.linear_c2 = MLP(input_dim=c2_in_channels, embed_dim=embedding_dim) | |
| self.linear_c1 = MLP(input_dim=c1_in_channels, embed_dim=embedding_dim) | |
| self.linear_fuse = nn.Sequential( | |
| nn.Conv2d(in_channels=embedding_dim*4, out_channels=embedding_dim, kernel_size=1), | |
| norm_layer(embedding_dim), | |
| nn.ReLU(inplace=True) | |
| ) | |
| self.linear_pred = nn.Conv2d(embedding_dim, self.num_classes, kernel_size=1) | |
| def forward(self, inputs, return_feats=False): | |
| # len=4, 1/4,1/8,1/16,1/32 | |
| c1, c2, c3, c4 = inputs | |
| ############## MLP decoder on C1-C4 ########### | |
| n, _, h, w = c4.shape | |
| _c4 = self.linear_c4(c4).permute(0,2,1).reshape(n, -1, c4.shape[2], c4.shape[3]) | |
| _c4 = F.interpolate(_c4, size=c1.size()[2:],mode='bilinear',align_corners=self.align_corners) | |
| _c3 = self.linear_c3(c3).permute(0,2,1).reshape(n, -1, c3.shape[2], c3.shape[3]) | |
| _c3 = F.interpolate(_c3, size=c1.size()[2:],mode='bilinear',align_corners=self.align_corners) | |
| _c2 = self.linear_c2(c2).permute(0,2,1).reshape(n, -1, c2.shape[2], c2.shape[3]) | |
| _c2 = F.interpolate(_c2, size=c1.size()[2:],mode='bilinear',align_corners=self.align_corners) | |
| _c1 = self.linear_c1(c1).permute(0,2,1).reshape(n, -1, c1.shape[2], c1.shape[3]) | |
| _c = torch.cat([_c4, _c3, _c2, _c1], dim=1) | |
| x = self.linear_fuse(_c) | |
| x = self.dropout(x) | |
| x = self.linear_pred(x) | |
| if return_feats: | |
| return x, _c | |
| else: | |
| return x | |