# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
from typing import List, Tuple, Type
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from .mlp import MLPBlock
class PromptEncoder(nn.Module):
def __init__(
self,
embed_dim: int,
image_embedding_size: Tuple[int, int],
input_image_size: Tuple[int, int],
) -> None:
"""
Encodes prompts for input to SAM's mask decoder.
Arguments:
embed_dim (int): The prompts' embedding dimension
image_embedding_size (tuple(int, int)): The spatial size of the
image embedding, as (H, W).
input_image_size (int): The padded size of the image as input
to the image encoder, as (H, W).
"""
super().__init__()
self.embed_dim = embed_dim
self.input_image_size = input_image_size
self.image_embedding_size = image_embedding_size
self.pe_layer = PositionEmbeddingRandom(embed_dim // 2)
self.invalid_points = nn.Embedding(1, embed_dim)
self.point_embeddings = nn.Embedding(1, embed_dim)
self.bbox_top_left_embeddings = nn.Embedding(1, embed_dim)
self.bbox_bottom_right_embeddings = nn.Embedding(1, embed_dim)
def get_dense_pe(self) -> torch.Tensor:
"""
Returns the positional encoding used to encode point prompts,
applied to a dense set of points the shape of the image encoding.
Returns:
torch.Tensor: Positional encoding with shape
1x(embed_dim)x(embedding_h)x(embedding_w)
"""
return self.pe_layer(self.image_embedding_size).unsqueeze(0)
def _embed_points(
self,
points: torch.Tensor,
labels: torch.Tensor,
) -> torch.Tensor:
"""Embeds point prompts."""
points = points + 0.5 # Shift to center of pixel
point_embedding = self.pe_layer.forward_with_coords(
points, self.input_image_size
)
invalid_label_ids = torch.eq(labels, -1)[:,:,None]
point_label_ids = torch.eq(labels, 1)[:,:,None]
topleft_label_ids = torch.eq(labels, 2)[:,:,None]
bottomright_label_ids = torch.eq(labels, 3)[:,:,None]
point_embedding = point_embedding + self.invalid_points.weight[:,None,:] * invalid_label_ids
point_embedding = point_embedding + self.point_embeddings.weight[:,None,:] * point_label_ids
point_embedding = point_embedding + self.bbox_top_left_embeddings.weight[:,None,:] * topleft_label_ids
point_embedding = point_embedding + self.bbox_bottom_right_embeddings.weight[:,None,:] * bottomright_label_ids
return point_embedding
def forward(
self,
coords,
labels,
) -> torch.Tensor:
"""
Embeds different types of prompts, returning both sparse and dense
embeddings.
Arguments:
points: A tensor of shape [B, 2]
labels: An integer tensor of shape [B] where each element is 1,2 or 3.
Returns:
torch.Tensor: sparse embeddings for the points and boxes, with shape
BxNx(embed_dim), where N is determined by the number of input points
and boxes.
"""
return self._embed_points(coords, labels)
class PositionEmbeddingRandom(nn.Module):
"""
Positional encoding using random spatial frequencies.
"""
def __init__(self, num_pos_feats: int) -> None:
super().__init__()
self.register_buffer(
"positional_encoding_gaussian_matrix", torch.randn((2, num_pos_feats))
)
def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:
"""Positionally encode points that are normalized to [0,1]."""
# assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
coords = 2 * coords - 1
coords = coords @ self.positional_encoding_gaussian_matrix
coords = 2 * np.pi * coords
# outputs d_1 x ... x d_n x C shape
return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)
def forward(self, size: Tuple[int, int]) -> torch.Tensor:
"""Generate positional encoding for a grid of the specified size."""
h, w = size
device = self.positional_encoding_gaussian_matrix.device
grid = torch.ones([h, w], device=device, dtype=torch.float32)
y_embed = grid.cumsum(dim=0) - 0.5
x_embed = grid.cumsum(dim=1) - 0.5
y_embed = y_embed / h
x_embed = x_embed / w
pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))
return pe.permute(2, 0, 1) # C x H x W
def forward_with_coords(
self, coords_input: torch.Tensor, image_size: Tuple[int, int]
) -> torch.Tensor:
"""Positionally encode points that are not normalized to [0,1]."""
coords = coords_input.clone()
coords[:, :, 0] = coords[:, :, 0] / image_size[1]
coords[:, :, 1] = coords[:, :, 1] / image_size[0]
return self._pe_encoding(coords.to(torch.float)) # B x N x C
class MaskDecoder(nn.Module):
def __init__(
self,
*,
transformer_dim: int,
transformer: nn.Module,
num_multimask_outputs: int,
activation: Type[nn.Module],
normalization_type: str,
normalize_before_activation: bool,
iou_head_depth: int,
iou_head_hidden_dim: int,
upscaling_layer_dims: List[int],
) -> None:
"""
Predicts masks given an image and prompt embeddings, using a
transformer architecture.
Arguments:
transformer_dim (int): the channel dimension of the transformer
transformer (nn.Module): the transformer used to predict masks
num_multimask_outputs (int): the number of masks to predict
when disambiguating masks
activation (nn.Module): the type of activation to use when
upscaling masks
iou_head_depth (int): the depth of the MLP used to predict
mask quality
iou_head_hidden_dim (int): the hidden dimension of the MLP
used to predict mask quality
"""
super().__init__()
self.transformer_dim = transformer_dim
self.transformer = transformer
self.num_multimask_outputs = num_multimask_outputs
self.iou_token = nn.Embedding(1, transformer_dim)
if num_multimask_outputs > 1:
self.num_mask_tokens = num_multimask_outputs + 1
else:
self.num_mask_tokens = 1
self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim)
output_dim_after_upscaling = transformer_dim
self.final_output_upscaling_layers = nn.ModuleList([])
for idx, layer_dims in enumerate(upscaling_layer_dims):
self.final_output_upscaling_layers.append(
nn.Sequential(
nn.ConvTranspose2d(
output_dim_after_upscaling,
layer_dims,
kernel_size=2,
stride=2,
),
nn.GroupNorm(1, layer_dims)
if idx < len(upscaling_layer_dims) - 1
else nn.Identity(),
activation(),
)
)
output_dim_after_upscaling = layer_dims
self.output_hypernetworks_mlps = nn.ModuleList(
[
MLPBlock(
input_dim=transformer_dim,
hidden_dim=transformer_dim,
output_dim=output_dim_after_upscaling,
num_layers=2,
act=activation,
)
for i in range(self.num_mask_tokens)
]
)
self.iou_prediction_head = MLPBlock(
input_dim=transformer_dim,
hidden_dim=iou_head_hidden_dim,
output_dim=self.num_mask_tokens,
num_layers=iou_head_depth,
act=activation,
)
def forward(
self,
image_embeddings: torch.Tensor,
image_pe: torch.Tensor,
sparse_prompt_embeddings: torch.Tensor,
multimask_output: bool,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Predict masks given image and prompt embeddings.
Arguments:
image_embeddings: A tensor of shape [B, C, H, W] or [B*max_num_queries, C, H, W]
image_pe (torch.Tensor): positional encoding with the shape of image_embeddings (the batch dimension is broadcastable).
sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes
multimask_output (bool): Whether to return multiple masks or a single
mask.
Returns:
torch.Tensor: batched predicted masks
torch.Tensor: batched predictions of mask quality
"""
(
batch_size,
max_num_queries,
sparse_embed_dim_1,
sparse_embed_dim_2,
) = sparse_prompt_embeddings.shape
(
_,
image_embed_dim_c,
image_embed_dim_h,
image_embed_dim_w,
) = image_embeddings.shape
# Tile the image embedding for all queries.
image_embeddings_tiled = torch.tile(
image_embeddings[:, None, :, :, :], [1, max_num_queries, 1, 1, 1]
).view(
batch_size * max_num_queries,
image_embed_dim_c,
image_embed_dim_h,
image_embed_dim_w,
)
sparse_prompt_embeddings = sparse_prompt_embeddings.reshape(
batch_size * max_num_queries, sparse_embed_dim_1, sparse_embed_dim_2
)
masks, iou_pred = self.predict_masks(
image_embeddings=image_embeddings_tiled,
image_pe=image_pe,
sparse_prompt_embeddings=sparse_prompt_embeddings,
)
if multimask_output and self.num_multimask_outputs > 1:
return masks[:, 1:, :], iou_pred[:, 1:]
else:
return masks[:, :1, :], iou_pred[:, :1]
def predict_masks(
self,
image_embeddings: torch.Tensor,
image_pe: torch.Tensor,
sparse_prompt_embeddings: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Predicts masks. See 'forward' for more details."""
# Concatenate output tokens
output_tokens = torch.cat(
[self.iou_token.weight, self.mask_tokens.weight], dim=0
)
output_tokens = output_tokens.unsqueeze(0).expand(
sparse_prompt_embeddings.size(0), -1, -1
)
tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1)
# Expand per-image data in batch direction to be per-mask
pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0)
b, c, h, w = image_embeddings.shape
hs, src = self.transformer(image_embeddings, pos_src, tokens)
iou_token_out = hs[:, 0, :]
mask_tokens_out = hs[:, 1 : (1 + self.num_mask_tokens), :]
# Upscale mask embeddings and predict masks using the mask tokens
upscaled_embedding = src.transpose(1, 2).view(b, c, h, w)
for upscaling_layer in self.final_output_upscaling_layers:
upscaled_embedding = upscaling_layer(upscaled_embedding)
hyper_in_list: List[torch.Tensor] = []
for i, output_hypernetworks_mlp in enumerate(self.output_hypernetworks_mlps):
hyper_in_list.append(output_hypernetworks_mlp(mask_tokens_out[:, i, :]))
hyper_in = torch.stack(hyper_in_list, dim=1)
b, c, h, w = upscaled_embedding.shape
masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w)
# Generate mask quality predictions
iou_pred = self.iou_prediction_head(iou_token_out)
return masks, iou_pred