File size: 4,187 Bytes
c6535db | 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 117 118 119 120 121 122 123 124 125 126 | # Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
"""Necks are the interface between a vision backbone and the rest of the detection model"""
from copy import deepcopy
from typing import List, Optional, Tuple
import torch
import torch.nn as nn
class Sam3DualViTDetNeck(nn.Module):
def __init__(
self,
trunk: nn.Module,
position_encoding: nn.Module,
d_model: int,
scale_factors=(4.0, 2.0, 1.0, 0.5),
add_sam2_neck: bool = False,
):
"""
SimpleFPN neck a la ViTDet
(From detectron2, very lightly adapted)
It supports a "dual neck" setting, where we have two identical necks (for SAM3 and SAM2), with different weights
:param trunk: the backbone
:param position_encoding: the positional encoding to use
:param d_model: the dimension of the model
"""
super().__init__()
self.trunk = trunk
self.position_encoding = position_encoding
self.convs = nn.ModuleList()
self.scale_factors = scale_factors
use_bias = True
dim: int = self.trunk.channel_list[-1]
for _, scale in enumerate(scale_factors):
current = nn.Sequential()
if scale == 4.0:
current.add_module(
"dconv_2x2_0",
nn.ConvTranspose2d(dim, dim // 2, kernel_size=2, stride=2),
)
current.add_module(
"gelu",
nn.GELU(),
)
current.add_module(
"dconv_2x2_1",
nn.ConvTranspose2d(dim // 2, dim // 4, kernel_size=2, stride=2),
)
out_dim = dim // 4
elif scale == 2.0:
current.add_module(
"dconv_2x2",
nn.ConvTranspose2d(dim, dim // 2, kernel_size=2, stride=2),
)
out_dim = dim // 2
elif scale == 1.0:
out_dim = dim
elif scale == 0.5:
current.add_module(
"maxpool_2x2",
nn.MaxPool2d(kernel_size=2, stride=2),
)
out_dim = dim
else:
raise NotImplementedError(f"scale_factor={scale} is not supported yet.")
current.add_module(
"conv_1x1",
nn.Conv2d(
in_channels=out_dim,
out_channels=d_model,
kernel_size=1,
bias=use_bias,
),
)
current.add_module(
"conv_3x3",
nn.Conv2d(
in_channels=d_model,
out_channels=d_model,
kernel_size=3,
padding=1,
bias=use_bias,
),
)
self.convs.append(current)
self.sam2_convs = None
if add_sam2_neck:
# Assumes sam2 neck is just a clone of the original neck
self.sam2_convs = deepcopy(self.convs)
def forward(
self, tensor_list: List[torch.Tensor]
) -> Tuple[
List[torch.Tensor],
List[torch.Tensor],
Optional[List[torch.Tensor]],
Optional[List[torch.Tensor]],
]:
xs = self.trunk(tensor_list)
sam3_out, sam3_pos = [], []
sam2_out, sam2_pos = None, None
if self.sam2_convs is not None:
sam2_out, sam2_pos = [], []
x = xs[-1] # simpleFPN
for i in range(len(self.convs)):
sam3_x_out = self.convs[i](x)
sam3_pos_out = self.position_encoding(sam3_x_out).to(sam3_x_out.dtype)
sam3_out.append(sam3_x_out)
sam3_pos.append(sam3_pos_out)
if self.sam2_convs is not None:
sam2_x_out = self.sam2_convs[i](x)
sam2_pos_out = self.position_encoding(sam2_x_out).to(sam2_x_out.dtype)
sam2_out.append(sam2_x_out)
sam2_pos.append(sam2_pos_out)
return sam3_out, sam3_pos, sam2_out, sam2_pos
|