Spaces:
Running
Running
File size: 4,891 Bytes
1b369eb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from ..utils.misc import NestedTensor, nested_tensor_from_tensor_list
import torchvision.transforms as transforms
from .backbone import build_backbone
from .deformable_transformer import build_deforamble_transformer
class BasicLayer(nn.Module):
"""
Basic Convolutional Layer: Conv2d -> BatchNorm -> ReLU
"""
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, bias=False):
super().__init__()
self.layer = nn.Sequential(
nn.Conv2d( in_channels, out_channels, kernel_size, padding = padding, stride=stride, dilation=dilation, bias = bias),
nn.BatchNorm2d(out_channels, affine=False),
nn.ReLU(inplace = False),
)
def forward(self, x):
return self.layer(x)
class RDD_Descriptor(nn.Module):
def __init__(self, backbone, transformer, num_feature_levels):
super().__init__()
self.transformer = transformer
self.hidden_dim = transformer.d_model
self.num_feature_levels = num_feature_levels
self.matchibility_head = nn.Sequential(
BasicLayer(256, 128, 1, padding=0),
BasicLayer(128, 64, 1, padding=0),
nn.Conv2d (64, 1, 1),
nn.Sigmoid()
)
if num_feature_levels > 1:
num_backbone_outs = len(backbone.strides)
input_proj_list = []
for _ in range(num_backbone_outs):
in_channels = backbone.num_channels[_]
input_proj_list.append(nn.Sequential(
nn.Conv2d(in_channels, self.hidden_dim, kernel_size=1),
nn.GroupNorm(32, self.hidden_dim),
))
for _ in range(num_feature_levels - num_backbone_outs):
input_proj_list.append(nn.Sequential(
nn.Conv2d(in_channels, self.hidden_dim, kernel_size=3, stride=2, padding=1),
nn.GroupNorm(32, self.hidden_dim),
))
in_channels = self.hidden_dim
self.input_proj = nn.ModuleList(input_proj_list)
else:
self.input_proj = nn.ModuleList([
nn.Sequential(
nn.Conv2d(backbone.num_channels[0], self.hidden_dim, kernel_size=1),
nn.GroupNorm(32, self.hidden_dim),
)])
self.backbone = backbone
self.stride = backbone.strides[0]
for proj in self.input_proj:
nn.init.xavier_uniform_(proj[0].weight, gain=1)
nn.init.constant_(proj[0].bias, 0)
def forward(self, samples: NestedTensor):
if not isinstance(samples, NestedTensor):
samples = nested_tensor_from_tensor_list(samples)
features, pos = self.backbone(samples)
srcs = []
masks = []
for l, feat in enumerate(features):
src, mask = feat.decompose()
srcs.append(self.input_proj[l](src))
masks.append(mask)
assert mask is not None
if self.num_feature_levels > len(srcs):
_len_srcs = len(srcs)
for l in range(_len_srcs, self.num_feature_levels):
if l == _len_srcs:
src = self.input_proj[l](features[-1].tensors)
else:
src = self.input_proj[l](srcs[-1])
m = samples.mask
mask = F.interpolate(m[None].float(), size=src.shape[-2:]).to(torch.bool)[0]
pos_l = self.backbone[1](NestedTensor(src, mask)).to(src.dtype)
srcs.append(src)
masks.append(mask)
pos.append(pos_l)
flatten_feats, spatial_shapes, level_start_index = self.transformer(srcs, masks, pos)
# Reshape the flattened features back to the original spatial shapes
feats = []
level_start_index = torch.cat((level_start_index, torch.tensor([flatten_feats.shape[1]+1]).to(level_start_index.device)))
for i, shape in enumerate(spatial_shapes):
assert len(shape) == 2
temp = flatten_feats[:, level_start_index[i] : level_start_index[i+1], :]
feats.append(temp.transpose(1, 2).view(-1, self.hidden_dim, *shape))
# Sum up the features from different levels
final_feature = feats[0]
for feat in feats[1:]:
final_feature = final_feature + F.interpolate(feat, size=final_feature.shape[-2:], mode='bilinear', align_corners=True)
matchibility = self.matchibility_head(final_feature)
return final_feature, matchibility
def build_descriptor(config):
backbone = build_backbone(config)
transformer = build_deforamble_transformer(config)
return RDD_Descriptor(backbone, transformer, config['num_feature_levels']) |