Spaces:
Runtime error
Runtime error
File size: 5,993 Bytes
9cc3eb2 |
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 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
from typing import Tuple, Literal
import torch
from mmengine import MMLogger
from mmdet.registry import MODELS
from mmengine.model import BaseModule
from mmengine.structures import InstanceData
from ext.sam import PromptEncoder
from ext.meta.sam_meta import meta_dict, checkpoint_dict
from utils.load_checkpoint import load_checkpoint_with_prefix
@MODELS.register_module()
class SAMPromptEncoder(BaseModule):
def __init__(
self,
model_name: Literal['vit_h', 'vit_l', 'vit_b'] = 'vit_h',
fix: bool = True,
init_cfg=None,
):
assert init_cfg is not None and init_cfg['type'] == 'sam_pretrain', f"{init_cfg['type']} is not supported."
pretrained = init_cfg['checkpoint']
super().__init__(init_cfg=None)
self.init_cfg = init_cfg
self.logger = MMLogger.get_current_instance()
backbone_meta = meta_dict[model_name]
checkpoint_path = checkpoint_dict[pretrained]
prompt_encoder = PromptEncoder(
embed_dim=256,
image_embedding_size=(backbone_meta['image_embedding_size'], backbone_meta['image_embedding_size']),
input_image_size=(backbone_meta['image_size'], backbone_meta['image_size']),
mask_in_chans=16,
)
state_dict = load_checkpoint_with_prefix(checkpoint_path, prefix='prompt_encoder')
prompt_encoder.load_state_dict(state_dict, strict=True)
# meta
self.embed_dim = prompt_encoder.embed_dim
self.input_image_size = prompt_encoder.input_image_size
self.image_embedding_size = prompt_encoder.image_embedding_size
self.num_point_embeddings = 4
self.mask_input_size = prompt_encoder.mask_input_size
# positional encoding
self.pe_layer = prompt_encoder.pe_layer
# mask encoding
self.mask_downscaling = prompt_encoder.mask_downscaling
self.no_mask_embed = prompt_encoder.no_mask_embed
# point encoding
self.point_embeddings = prompt_encoder.point_embeddings
self.not_a_point_embed = prompt_encoder.not_a_point_embed
self.fix = fix
if self.fix:
self.train(mode=False)
for name, param in self.named_parameters():
param.requires_grad = False
@property
def device(self):
return self.no_mask_embed.weight.device
def init_weights(self):
self.logger.info(f"Init Config for {self.__class__.__name__}")
self.logger.info(self.init_cfg)
def train(self: torch.nn.Module, mode: bool = True) -> torch.nn.Module:
if not isinstance(mode, bool):
raise ValueError("training mode is expected to be boolean")
if self.fix:
super().train(mode=False)
else:
super().train(mode=mode)
return self
def _embed_boxes(self, bboxes: torch.Tensor, image_size: Tuple[int, int]) -> torch.Tensor:
"""Embeds box prompts."""
bboxes = bboxes + 0.5 # Shift to center of pixel
coords = bboxes.reshape(-1, 2, 2)
corner_embedding = self.pe_layer.forward_with_coords(coords, image_size)
corner_embedding[:, 0, :] += self.point_embeddings[2].weight
corner_embedding[:, 1, :] += self.point_embeddings[3].weight
return corner_embedding
def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor:
"""Embeds mask inputs."""
mask_embedding = self.mask_downscaling(masks)
return mask_embedding
def get_dense_pe(self) -> torch.Tensor:
return self.pe_layer(self.image_embedding_size).unsqueeze(0)
def _embed_points(
self,
points: torch.Tensor,
labels: torch.Tensor,
pad: bool,
) -> torch.Tensor:
"""Embeds point prompts."""
points = points + 0.5 # Shift to center of pixel
if pad:
padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device)
padding_label = -torch.ones((labels.shape[0], 1), device=labels.device)
points = torch.cat([points, padding_point], dim=1)
labels = torch.cat([labels, padding_label], dim=1)
point_embedding = self.pe_layer.forward_with_coords(points, self.input_image_size)
point_embedding[labels == -1] = 0.0
point_embedding[labels == -1] += self.not_a_point_embed.weight
point_embedding[labels == 0] += self.point_embeddings[0].weight
point_embedding[labels == 1] += self.point_embeddings[1].weight
return point_embedding
def forward(
self,
instances: InstanceData,
image_size: Tuple[int, int],
with_points: bool = False,
with_bboxes: bool = False,
with_masks: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor]:
assert with_points or with_bboxes or with_masks
bs = len(instances)
sparse_embeddings = torch.empty((bs, 0, self.embed_dim), device=self.device)
if with_points:
assert 'point_coords' in instances
coords = instances.point_coords
labels = torch.ones_like(coords)[:, :, 0]
point_embeddings = self._embed_points(coords, labels, pad=not with_bboxes)
sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
if with_bboxes:
assert 'bboxes' in instances
box_embeddings = self._embed_boxes(
instances.bboxes, image_size=image_size
)
sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)
if with_masks:
assert 'masks' in instances
dense_embeddings = self._embed_masks(instances.masks.masks)
else:
dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand(
bs, -1, self.image_embedding_size[0], self.image_embedding_size[1]
)
return sparse_embeddings, dense_embeddings
|