| import torch.nn.functional as F |
| from mmengine.model import BaseModel |
|
|
| from mmdet.registry import MODELS |
| from mmdet.utils import ConfigType, OptConfigType, OptMultiConfig |
|
|
|
|
| @MODELS.register_module() |
| class SAMSegmentor(BaseModel): |
| MASK_THRESHOLD = 0.5 |
|
|
| def __init__( |
| self, |
| backbone: ConfigType, |
| neck: ConfigType, |
| prompt_encoder: ConfigType, |
| mask_decoder: ConfigType, |
| data_preprocessor: OptConfigType = None, |
| fpn_neck: OptConfigType = None, |
| init_cfg: OptMultiConfig = None, |
| use_clip_feat: bool = False, |
| use_head_feat: bool = False, |
| use_gt_prompt: bool = False, |
| use_point: bool = False, |
| enable_backbone: bool = False, |
| ) -> None: |
| super().__init__(data_preprocessor=data_preprocessor, init_cfg=init_cfg) |
|
|
| self.backbone = MODELS.build(backbone) |
| self.neck = MODELS.build(neck) |
| self.pe = MODELS.build(prompt_encoder) |
| self.mask_decoder = MODELS.build(mask_decoder) |
| if fpn_neck is not None: |
| self.fpn_neck = MODELS.build(fpn_neck) |
| else: |
| self.fpn_neck = None |
|
|
| self.use_clip_feat = use_clip_feat |
| self.use_head_feat = use_head_feat |
| self.use_gt_prompt = use_gt_prompt |
| self.use_point = use_point |
|
|
| self.enable_backbone = enable_backbone |
|
|
| def extract_feat(self, inputs): |
| backbone_feat = self.backbone(inputs) |
| neck_feat = self.neck(backbone_feat) |
| if self.fpn_neck is not None: |
| fpn_feat = self.fpn_neck(backbone_feat) |
| else: |
| fpn_feat = None |
|
|
| return dict( |
| backbone_feat=backbone_feat, |
| neck_feat=neck_feat, |
| fpn_feat=fpn_feat |
| ) |
|
|
| def extract_masks(self, feat_cache, prompts): |
| sparse_embed, dense_embed = self.pe( |
| prompts, |
| image_size=(1024, 1024), |
| with_points='point_coords' in prompts, |
| with_bboxes='bboxes' in prompts, |
| ) |
|
|
| kwargs = dict() |
| if self.enable_backbone: |
| kwargs['backbone_feats'] = feat_cache['backbone_feat'] |
| kwargs['backbone'] = self.backbone |
| kwargs['fpn_feats'] = feat_cache['fpn_feat'] |
| low_res_masks, iou_predictions, cls_pred = self.mask_decoder( |
| image_embeddings=feat_cache['neck_feat'], |
| image_pe=self.pe.get_dense_pe(), |
| sparse_prompt_embeddings=sparse_embed, |
| dense_prompt_embeddings=dense_embed, |
| multi_mask_output=False, |
| **kwargs |
| ) |
| masks = F.interpolate( |
| low_res_masks, |
| scale_factor=4., |
| mode='bilinear', |
| align_corners=False, |
| ) |
|
|
| masks = masks.sigmoid() |
| cls_pred = cls_pred.softmax(-1)[..., :-1] |
| return masks.detach().cpu().numpy(), cls_pred.detach().cpu() |
|
|
| def forward(self, inputs): |
| return inputs |
|
|
