| |
|
|
| """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: |
| |
| 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] |
| 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 |
|
|
