|
|
import numpy as np |
|
|
import torch.nn as nn |
|
|
import torch |
|
|
|
|
|
from torch.nn.modules import module |
|
|
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): |
|
|
|
|
|
c1, c2, c3, c4 = inputs |
|
|
|
|
|
|
|
|
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 |
|
|
|
