Bỏ để qua phần nội dung

Tài liệu tham khảo cho ultralytics/models/sam/modules/decoders.py

Ghi

Tệp này có sẵn tại https://github.com/ultralytics/ultralytics/blob/main/ultralytics/Mô hình/sam/mô-đun/decoders.py. Nếu bạn phát hiện ra một vấn đề, vui lòng giúp khắc phục nó bằng cách đóng góp Yêu cầu 🛠️ kéo. Cảm ơn bạn 🙏 !



ultralytics.models.sam.modules.decoders.MaskDecoder

Căn cứ: Module

Mô-đun giải mã để tạo mặt nạ và điểm chất lượng liên quan của chúng, sử dụng kiến trúc máy biến áp để dự đoán mặt nạ được đưa ra hình ảnh và nhúng nhanh chóng.

Thuộc tính:

Tên Kiểu Sự miêu tả
transformer_dim int

Kích thước kênh cho mô-đun biến áp.

transformer Module

Mô-đun biến áp được sử dụng để dự đoán mặt nạ.

num_multimask_outputs int

Số lượng mặt nạ để dự đoán cho mặt nạ định hướng.

iou_token Embedding

Nhúng cho mã thông báo IoU.

num_mask_tokens int

Số lượng mã thông báo mặt nạ.

mask_tokens Embedding

Nhúng cho mã thông báo mặt nạ.

output_upscaling Sequential

Trình tự mạng nơ-ron để nâng cấp đầu ra.

output_hypernetworks_mlps ModuleList

MLP siêu mạng để tạo mặt nạ.

iou_prediction_head Module

MLP để dự đoán chất lượng khẩu trang.

Mã nguồn trong ultralytics/models/sam/modules/decoders.py
class MaskDecoder(nn.Module):
    """
    Decoder module for generating masks and their associated quality scores, using a transformer architecture to predict
    masks given image and prompt embeddings.

    Attributes:
        transformer_dim (int): Channel dimension for the transformer module.
        transformer (nn.Module): The transformer module used for mask prediction.
        num_multimask_outputs (int): Number of masks to predict for disambiguating masks.
        iou_token (nn.Embedding): Embedding for the IoU token.
        num_mask_tokens (int): Number of mask tokens.
        mask_tokens (nn.Embedding): Embedding for the mask tokens.
        output_upscaling (nn.Sequential): Neural network sequence for upscaling the output.
        output_hypernetworks_mlps (nn.ModuleList): Hypernetwork MLPs for generating masks.
        iou_prediction_head (nn.Module): MLP for predicting mask quality.
    """

    def __init__(
        self,
        *,
        transformer_dim: int,
        transformer: nn.Module,
        num_multimask_outputs: int = 3,
        activation: Type[nn.Module] = nn.GELU,
        iou_head_depth: int = 3,
        iou_head_hidden_dim: int = 256,
    ) -> None:
        """
        Predicts masks given an image and prompt embeddings, using a transformer architecture.

        Args:
            transformer_dim (int): the channel dimension of the transformer module
            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)
        self.num_mask_tokens = num_multimask_outputs + 1
        self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim)

        self.output_upscaling = nn.Sequential(
            nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2),
            LayerNorm2d(transformer_dim // 4),
            activation(),
            nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2),
            activation(),
        )
        self.output_hypernetworks_mlps = nn.ModuleList(
            [MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3) for _ in range(self.num_mask_tokens)]
        )

        self.iou_prediction_head = MLP(transformer_dim, iou_head_hidden_dim, self.num_mask_tokens, iou_head_depth)

    def forward(
        self,
        image_embeddings: torch.Tensor,
        image_pe: torch.Tensor,
        sparse_prompt_embeddings: torch.Tensor,
        dense_prompt_embeddings: torch.Tensor,
        multimask_output: bool,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Predict masks given image and prompt embeddings.

        Args:
            image_embeddings (torch.Tensor): the embeddings from the image encoder
            image_pe (torch.Tensor): positional encoding with the shape of image_embeddings
            sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes
            dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs
            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
        """
        masks, iou_pred = self.predict_masks(
            image_embeddings=image_embeddings,
            image_pe=image_pe,
            sparse_prompt_embeddings=sparse_prompt_embeddings,
            dense_prompt_embeddings=dense_prompt_embeddings,
        )

        # Select the correct mask or masks for output
        mask_slice = slice(1, None) if multimask_output else slice(0, 1)
        masks = masks[:, mask_slice, :, :]
        iou_pred = iou_pred[:, mask_slice]

        # Prepare output
        return masks, iou_pred

    def predict_masks(
        self,
        image_embeddings: torch.Tensor,
        image_pe: torch.Tensor,
        sparse_prompt_embeddings: torch.Tensor,
        dense_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.shape[0], -1, -1)
        tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1)

        # Expand per-image data in batch direction to be per-mask
        src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0)
        src = src + dense_prompt_embeddings
        pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0)
        b, c, h, w = src.shape

        # Run the transformer
        hs, src = self.transformer(src, 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
        src = src.transpose(1, 2).view(b, c, h, w)
        upscaled_embedding = self.output_upscaling(src)
        hyper_in_list: List[torch.Tensor] = [
            self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :]) for i in range(self.num_mask_tokens)
        ]
        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

__init__(*, transformer_dim, transformer, num_multimask_outputs=3, activation=nn.GELU, iou_head_depth=3, iou_head_hidden_dim=256)

Dự đoán mặt nạ được cung cấp hình ảnh và nhúng nhanh chóng, sử dụng kiến trúc biến áp.

Thông số:

Tên Kiểu Sự miêu tả Mặc định
transformer_dim int

Kích thước kênh của mô-đun máy biến áp

bắt buộc
transformer Module

Máy biến áp được sử dụng để dự đoán mặt nạ

bắt buộc
num_multimask_outputs int

Số lượng mặt nạ cần dự đoán khi định hướng mặt nạ

3
activation Module

Loại kích hoạt để sử dụng khi nâng cấp mặt nạ

GELU
iou_head_depth int

độ sâu của MLP được sử dụng để dự đoán chất lượng mặt nạ

3
iou_head_hidden_dim int

kích thước ẩn của MLP được sử dụng để dự đoán chất lượng mặt nạ

256
Mã nguồn trong ultralytics/models/sam/modules/decoders.py
def __init__(
    self,
    *,
    transformer_dim: int,
    transformer: nn.Module,
    num_multimask_outputs: int = 3,
    activation: Type[nn.Module] = nn.GELU,
    iou_head_depth: int = 3,
    iou_head_hidden_dim: int = 256,
) -> None:
    """
    Predicts masks given an image and prompt embeddings, using a transformer architecture.

    Args:
        transformer_dim (int): the channel dimension of the transformer module
        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)
    self.num_mask_tokens = num_multimask_outputs + 1
    self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim)

    self.output_upscaling = nn.Sequential(
        nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2),
        LayerNorm2d(transformer_dim // 4),
        activation(),
        nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2),
        activation(),
    )
    self.output_hypernetworks_mlps = nn.ModuleList(
        [MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3) for _ in range(self.num_mask_tokens)]
    )

    self.iou_prediction_head = MLP(transformer_dim, iou_head_hidden_dim, self.num_mask_tokens, iou_head_depth)

forward(image_embeddings, image_pe, sparse_prompt_embeddings, dense_prompt_embeddings, multimask_output)

Dự đoán mặt nạ đã cho hình ảnh và nhúng nhanh chóng.

Thông số:

Tên Kiểu Sự miêu tả Mặc định
image_embeddings Tensor

Các nhúng từ bộ mã hóa hình ảnh

bắt buộc
image_pe Tensor

Mã hóa vị trí với hình dạng của image_embeddings

bắt buộc
sparse_prompt_embeddings Tensor

Việc nhúng các điểm và hộp

bắt buộc
dense_prompt_embeddings Tensor

Việc nhúng các đầu vào mặt nạ

bắt buộc
multimask_output bool

Có nên trả lại nhiều khẩu trang hay một khẩu trang duy nhất.

bắt buộc

Trở lại:

Kiểu Sự miêu tả
Tensor

torch.Tensor: Mặt nạ dự đoán hàng loạt

Tensor

torch.Tensor: Dự đoán hàng loạt về chất lượng mặt nạ

Mã nguồn trong ultralytics/models/sam/modules/decoders.py
def forward(
    self,
    image_embeddings: torch.Tensor,
    image_pe: torch.Tensor,
    sparse_prompt_embeddings: torch.Tensor,
    dense_prompt_embeddings: torch.Tensor,
    multimask_output: bool,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Predict masks given image and prompt embeddings.

    Args:
        image_embeddings (torch.Tensor): the embeddings from the image encoder
        image_pe (torch.Tensor): positional encoding with the shape of image_embeddings
        sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes
        dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs
        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
    """
    masks, iou_pred = self.predict_masks(
        image_embeddings=image_embeddings,
        image_pe=image_pe,
        sparse_prompt_embeddings=sparse_prompt_embeddings,
        dense_prompt_embeddings=dense_prompt_embeddings,
    )

    # Select the correct mask or masks for output
    mask_slice = slice(1, None) if multimask_output else slice(0, 1)
    masks = masks[:, mask_slice, :, :]
    iou_pred = iou_pred[:, mask_slice]

    # Prepare output
    return masks, iou_pred

predict_masks(image_embeddings, image_pe, sparse_prompt_embeddings, dense_prompt_embeddings)

Dự đoán khẩu trang.

Xem 'chuyển tiếp' để biết thêm chi tiết.

Mã nguồn trong ultralytics/models/sam/modules/decoders.py
def predict_masks(
    self,
    image_embeddings: torch.Tensor,
    image_pe: torch.Tensor,
    sparse_prompt_embeddings: torch.Tensor,
    dense_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.shape[0], -1, -1)
    tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1)

    # Expand per-image data in batch direction to be per-mask
    src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0)
    src = src + dense_prompt_embeddings
    pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0)
    b, c, h, w = src.shape

    # Run the transformer
    hs, src = self.transformer(src, 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
    src = src.transpose(1, 2).view(b, c, h, w)
    upscaled_embedding = self.output_upscaling(src)
    hyper_in_list: List[torch.Tensor] = [
        self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :]) for i in range(self.num_mask_tokens)
    ]
    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



ultralytics.models.sam.modules.decoders.MLP

Căn cứ: Module

Mô hình MLP (Multi-Layer Perceptron) được điều chỉnh nhẹ từ https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py

Mã nguồn trong ultralytics/models/sam/modules/decoders.py
class MLP(nn.Module):
    """
    MLP (Multi-Layer Perceptron) model lightly adapted from
    https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py
    """

    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        output_dim: int,
        num_layers: int,
        sigmoid_output: bool = False,
    ) -> None:
        """
        Initializes the MLP (Multi-Layer Perceptron) model.

        Args:
            input_dim (int): The dimensionality of the input features.
            hidden_dim (int): The dimensionality of the hidden layers.
            output_dim (int): The dimensionality of the output layer.
            num_layers (int): The number of hidden layers.
            sigmoid_output (bool, optional): Apply a sigmoid activation to the output layer. Defaults to False.
        """
        super().__init__()
        self.num_layers = num_layers
        h = [hidden_dim] * (num_layers - 1)
        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
        self.sigmoid_output = sigmoid_output

    def forward(self, x):
        """Executes feedforward within the neural network module and applies activation."""
        for i, layer in enumerate(self.layers):
            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
        if self.sigmoid_output:
            x = torch.sigmoid(x)
        return x

__init__(input_dim, hidden_dim, output_dim, num_layers, sigmoid_output=False)

Khởi tạo mô hình MLP (Multi-Layer Perceptron).

Thông số:

Tên Kiểu Sự miêu tả Mặc định
input_dim int

Tính kích thước của các tính năng đầu vào.

bắt buộc
hidden_dim int

Tính chiều của các lớp ẩn.

bắt buộc
output_dim int

Tính kích thước của lớp đầu ra.

bắt buộc
num_layers int

Số lượng các lớp ẩn.

bắt buộc
sigmoid_output bool

Áp dụng kích hoạt sigmoid cho lớp đầu ra. Mặc định là False.

False
Mã nguồn trong ultralytics/models/sam/modules/decoders.py
def __init__(
    self,
    input_dim: int,
    hidden_dim: int,
    output_dim: int,
    num_layers: int,
    sigmoid_output: bool = False,
) -> None:
    """
    Initializes the MLP (Multi-Layer Perceptron) model.

    Args:
        input_dim (int): The dimensionality of the input features.
        hidden_dim (int): The dimensionality of the hidden layers.
        output_dim (int): The dimensionality of the output layer.
        num_layers (int): The number of hidden layers.
        sigmoid_output (bool, optional): Apply a sigmoid activation to the output layer. Defaults to False.
    """
    super().__init__()
    self.num_layers = num_layers
    h = [hidden_dim] * (num_layers - 1)
    self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
    self.sigmoid_output = sigmoid_output

forward(x)

Thực thi feedforward trong mô-đun mạng nơ-ron và áp dụng kích hoạt.

Mã nguồn trong ultralytics/models/sam/modules/decoders.py
def forward(self, x):
    """Executes feedforward within the neural network module and applies activation."""
    for i, layer in enumerate(self.layers):
        x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
    if self.sigmoid_output:
        x = torch.sigmoid(x)
    return x





Created 2023-11-12, Updated 2024-06-02
Authors: glenn-jocher (5), Burhan-Q (1), Laughing-q (1)