# Copyright 2022 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """PyTorch MaskFormer model.""" from dataclasses import dataclass from numbers import Number import numpy as np import torch from huggingface_hub.dataclasses import strict from torch import Tensor, nn from ... import initialization as init from ...configuration_utils import PreTrainedConfig, SubConfigSpec from ...modeling_utils import PreTrainedModel from ...utils import ( ModelOutput, auto_docstring, is_accelerate_available, is_scipy_available, logging, requires_backends, ) from ..auto import AutoBackbone, AutoConfig from ..detr.configuration_detr import DetrConfig from ..detr.modeling_detr import DetrDecoder, DetrDecoderOutput, DetrSinePositionEmbedding from .configuration_maskformer_swin import MaskFormerSwinConfig if is_accelerate_available(): from accelerate import PartialState from accelerate.utils import reduce if is_scipy_available(): from scipy.optimize import linear_sum_assignment logger = logging.get_logger(__name__) @auto_docstring(checkpoint="facebook/maskformer-swin-base-ade") @strict class MaskFormerDetrConfig(DetrConfig): model_type = "detr" @auto_docstring(checkpoint="facebook/maskformer-swin-base-ade") @strict class MaskFormerConfig(PreTrainedConfig): r""" fpn_feature_size (`int`, *optional*, defaults to 256): The Feature Pyramid Network's features size. mask_feature_size (`int`, *optional*, defaults to 256): The masks' features size, this value will also be used to specify the Feature Pyramid Network features' size. decoder_config (`Dict`, *optional*): The configuration passed to the transformer decoder model, if unset the base config for `detr-resnet-50` will be used. cross_entropy_weight (`float`, *optional*, defaults to 1.0): The weight for the cross entropy loss. output_auxiliary_logits (`bool`, *optional*): Should the model output its `auxiliary_logits` or not. Raises: `ValueError`: Raised if the backbone model type selected is not in `["swin"]` or the decoder model type selected is not in `["detr"]` Examples: ```python >>> from transformers import MaskFormerConfig, MaskFormerModel >>> # Initializing a MaskFormer facebook/maskformer-swin-base-ade configuration >>> configuration = MaskFormerConfig() >>> # Initializing a model (with random weights) from the facebook/maskformer-swin-base-ade style configuration >>> model = MaskFormerModel(configuration) >>> # Accessing the model configuration >>> configuration = model.config ``` """ model_type = "maskformer" sub_configs_defaults = { "backbone_config": SubConfigSpec( config_class=AutoConfig, model_type="swin", init_kwargs={ "depths": [2, 2, 18, 2], "drop_path_rate": 0.3, "image_size": 384, "embed_dim": 128, "num_heads": [4, 8, 16, 32], "window_size": 12, "out_features": ["stage1", "stage2", "stage3", "stage4"], }, ), "decoder_config": SubConfigSpec(config_class=AutoConfig, model_type="detr"), } attribute_map = {"hidden_size": "mask_feature_size"} backbones_supported = ["resnet", "swin"] decoders_supported = ["detr"] fpn_feature_size: int = 256 mask_feature_size: int = 256 no_object_weight: float = 0.1 use_auxiliary_loss: bool = False backbone_config: dict | PreTrainedConfig | None = None decoder_config: dict | PreTrainedConfig | None = None init_std: float = 0.02 init_xavier_std: float = 1.0 dice_weight: float = 1.0 cross_entropy_weight: float = 1.0 mask_weight: float = 20.0 output_auxiliary_logits: bool | None = None def __post_init__(self, **kwargs): super().__post_init__(**kwargs) self.num_attention_heads = self.decoder_config.encoder_attention_heads self.num_hidden_layers = self.decoder_config.num_hidden_layers def validate_architecture(self): super().validate_architecture() if self.backbone_config.model_type not in self.backbones_supported: logger.warning_once( f"Backbone {self.backbone_config.model_type} is not a supported model and may not be compatible with Mask2Former. " f"Supported model types: {','.join(self.backbones_supported)}" ) if self.decoder_config.model_type not in self.decoders_supported: raise ValueError( f"Transformer Decoder {self.decoder_config.model_type} not supported, please use one of" f" {','.join(self.decoders_supported)}" ) class DetrDecoderOutput(DetrDecoderOutput): pass @auto_docstring( custom_intro=""" MaskFormer's pixel level module output. It returns both the last and (optionally) the hidden states from the `encoder` and `decoder`. By default, the `encoder` is a MaskFormerSwin Transformer and the `decoder` is a Feature Pyramid Network (FPN). The `encoder_last_hidden_state` are referred on the paper as **images features**, while `decoder_last_hidden_state` as **pixel embeddings** """ ) @dataclass class MaskFormerPixelLevelModuleOutput(ModelOutput): r""" encoder_last_hidden_state (`torch.FloatTensor` of shape`(batch_size, num_channels, height, width)`): Last hidden states (final feature map) of the last stage of the encoder. decoder_last_hidden_state (`torch.FloatTensor` of shape`(batch_size, num_channels, height, width)`): Last hidden states (final feature map) of the last stage of the decoder. encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the model at the output of each stage. decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the model at the output of each stage. """ encoder_last_hidden_state: torch.FloatTensor | None = None decoder_last_hidden_state: torch.FloatTensor | None = None encoder_hidden_states: tuple[torch.FloatTensor] | None = None decoder_hidden_states: tuple[torch.FloatTensor] | None = None @auto_docstring( custom_intro=""" MaskFormer's pixel decoder module output, practically a Feature Pyramid Network. It returns the last hidden state and (optionally) the hidden states. """ ) @dataclass class MaskFormerPixelDecoderOutput(ModelOutput): r""" last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Last hidden states (final feature map) of the last stage of the model. """ last_hidden_state: torch.FloatTensor | None = None hidden_states: tuple[torch.FloatTensor] | None = None attentions: tuple[torch.FloatTensor] | None = None @auto_docstring( custom_intro=""" Class for outputs of [`MaskFormerModel`]. This class returns all the needed hidden states to compute the logits. """ ) @dataclass class MaskFormerModelOutput(ModelOutput): r""" encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Last hidden states (final feature map) of the last stage of the encoder model (backbone). pixel_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Last hidden states (final feature map) of the last stage of the pixel decoder model (FPN). transformer_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Last hidden states (final feature map) of the last stage of the transformer decoder model. encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the encoder model at the output of each stage. pixel_decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the pixel decoder model at the output of each stage. transformer_decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states (also called feature maps) of the transformer decoder at the output of each stage. hidden_states `tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` containing `encoder_hidden_states`, `pixel_decoder_hidden_states` and `decoder_hidden_states` """ encoder_last_hidden_state: torch.FloatTensor | None = None pixel_decoder_last_hidden_state: torch.FloatTensor | None = None transformer_decoder_last_hidden_state: torch.FloatTensor | None = None encoder_hidden_states: tuple[torch.FloatTensor] | None = None pixel_decoder_hidden_states: tuple[torch.FloatTensor] | None = None transformer_decoder_hidden_states: tuple[torch.FloatTensor] | None = None hidden_states: tuple[torch.FloatTensor] | None = None attentions: tuple[torch.FloatTensor] | None = None @auto_docstring( custom_intro=""" Class for outputs of [`MaskFormerForInstanceSegmentation`]. This output can be directly passed to [`~MaskFormerImageProcessor.post_process_semantic_segmentation`] or [`~MaskFormerImageProcessor.post_process_instance_segmentation`] or [`~MaskFormerImageProcessor.post_process_panoptic_segmentation`] depending on the task. Please, see [`~MaskFormerImageProcessor] for details regarding usage. """ ) @dataclass class MaskFormerForInstanceSegmentationOutput(ModelOutput): r""" loss (`torch.Tensor`, *optional*): The computed loss, returned when labels are present. class_queries_logits (`torch.FloatTensor`): A tensor of shape `(batch_size, num_queries, num_labels + 1)` representing the proposed classes for each query. Note the `+ 1` is needed because we incorporate the null class. masks_queries_logits (`torch.FloatTensor`): A tensor of shape `(batch_size, num_queries, height, width)` representing the proposed masks for each query. auxiliary_logits (`Dict[str, torch.FloatTensor]`, *optional*, returned when `output_auxiliary_logits=True`): Dictionary containing auxiliary predictions for each decoder layer when auxiliary losses are enabled. encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Last hidden states (final feature map) of the last stage of the encoder model (backbone). pixel_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Last hidden states (final feature map) of the last stage of the pixel decoder model (FPN). transformer_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Last hidden states (final feature map) of the last stage of the transformer decoder model. encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the encoder model at the output of each stage. pixel_decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the pixel decoder model at the output of each stage. transformer_decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the transformer decoder at the output of each stage. hidden_states `tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` containing `encoder_hidden_states`, `pixel_decoder_hidden_states` and `decoder_hidden_states`. """ loss: torch.FloatTensor | None = None class_queries_logits: torch.FloatTensor | None = None masks_queries_logits: torch.FloatTensor | None = None auxiliary_logits: torch.FloatTensor | None = None encoder_last_hidden_state: torch.FloatTensor | None = None pixel_decoder_last_hidden_state: torch.FloatTensor | None = None transformer_decoder_last_hidden_state: torch.FloatTensor | None = None encoder_hidden_states: tuple[torch.FloatTensor] | None = None pixel_decoder_hidden_states: tuple[torch.FloatTensor] | None = None transformer_decoder_hidden_states: tuple[torch.FloatTensor] | None = None hidden_states: tuple[torch.FloatTensor] | None = None attentions: tuple[torch.FloatTensor] | None = None def upsample_like(pixel_values: Tensor, like: Tensor, mode: str = "bilinear") -> Tensor: """ An utility function that upsamples `pixel_values` to match the dimension of `like`. Args: pixel_values (`torch.Tensor`): The tensor we wish to upsample. like (`torch.Tensor`): The tensor we wish to use as size target. mode (str, *optional*, defaults to `"bilinear"`): The interpolation mode. Returns: `torch.Tensor`: The upsampled tensor """ _, _, height, width = like.shape upsampled = nn.functional.interpolate(pixel_values, size=(height, width), mode=mode, align_corners=False) return upsampled # refactored from original implementation def dice_loss(inputs: Tensor, labels: Tensor, num_masks: int) -> Tensor: r""" Compute the DICE loss, similar to generalized IOU for masks as follows: $$ \mathcal{L}_{\text{dice}(x, y) = 1 - \frac{2 * x \cap y }{x \cup y + 1}} $$ In practice, since `labels` is a binary mask, (only 0s and 1s), dice can be computed as follow $$ \mathcal{L}_{\text{dice}(x, y) = 1 - \frac{2 * x * y }{x + y + 1}} $$ Args: inputs (`torch.Tensor`): A tensor representing a mask. labels (`torch.Tensor`): A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs (0 for the negative class and 1 for the positive class). num_masks (`int`): The number of masks present in the current batch, used for normalization. Returns: `torch.Tensor`: The computed loss. """ probs = inputs.sigmoid().flatten(1) numerator = 2 * (probs * labels).sum(-1) denominator = probs.sum(-1) + labels.sum(-1) loss = 1 - (numerator + 1) / (denominator + 1) loss = loss.sum() / num_masks return loss # refactored from original implementation def sigmoid_focal_loss( inputs: Tensor, labels: Tensor, num_masks: int, alpha: float = 0.25, gamma: float = 2 ) -> Tensor: r""" Focal loss proposed in [Focal Loss for Dense Object Detection](https://huggingface.co/papers/1708.02002) originally used in RetinaNet. The loss is computed as follows: $$ \mathcal{L}_{\text{focal loss} = -(1 - p_t)^{\gamma}\log{(p_t)} $$ where \\(CE(p_t) = -\log{(p_t)}}\\), CE is the standard Cross Entropy Loss Please refer to equation (1,2,3) of the paper for a better understanding. Args: inputs (`torch.Tensor`): A float tensor of arbitrary shape. labels (`torch.Tensor`): A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs (0 for the negative class and 1 for the positive class). num_masks (`int`): The number of masks present in the current batch, used for normalization. alpha (float, *optional*, defaults to 0.25): Weighting factor in range (0,1) to balance positive vs negative examples. gamma (float, *optional*, defaults to 2.0): Exponent of the modulating factor \\(1 - p_t\\) to balance easy vs hard examples. Returns: `torch.Tensor`: The computed loss. """ criterion = nn.BCEWithLogitsLoss(reduction="none") probs = inputs.sigmoid() cross_entropy_loss = criterion(inputs, labels) p_t = probs * labels + (1 - probs) * (1 - labels) loss = cross_entropy_loss * ((1 - p_t) ** gamma) if alpha >= 0: alpha_t = alpha * labels + (1 - alpha) * (1 - labels) loss = alpha_t * loss loss = loss.mean(1).sum() / num_masks return loss # refactored from original implementation def pair_wise_dice_loss(inputs: Tensor, labels: Tensor) -> Tensor: """ A pair wise version of the dice loss, see `dice_loss` for usage. Args: inputs (`torch.Tensor`): A tensor representing a mask labels (`torch.Tensor`): A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs (0 for the negative class and 1 for the positive class). Returns: `torch.Tensor`: The computed loss between each pairs. """ inputs = inputs.sigmoid().flatten(1) numerator = 2 * torch.matmul(inputs, labels.T) # using broadcasting to get a [num_queries, NUM_CLASSES] matrix denominator = inputs.sum(-1)[:, None] + labels.sum(-1)[None, :] loss = 1 - (numerator + 1) / (denominator + 1) return loss # refactored from original implementation def pair_wise_sigmoid_focal_loss(inputs: Tensor, labels: Tensor, alpha: float = 0.25, gamma: float = 2.0) -> Tensor: r""" A pair wise version of the focal loss, see `sigmoid_focal_loss` for usage. Args: inputs (`torch.Tensor`): A tensor representing a mask. labels (`torch.Tensor`): A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs (0 for the negative class and 1 for the positive class). alpha (float, *optional*, defaults to 0.25): Weighting factor in range (0,1) to balance positive vs negative examples. gamma (float, *optional*, defaults to 2.0): Exponent of the modulating factor \\(1 - p_t\\) to balance easy vs hard examples. Returns: `torch.Tensor`: The computed loss between each pairs. """ if alpha < 0: raise ValueError("alpha must be positive") height_and_width = inputs.shape[1] criterion = nn.BCEWithLogitsLoss(reduction="none") prob = inputs.sigmoid() cross_entropy_loss_pos = criterion(inputs, torch.ones_like(inputs)) focal_pos = ((1 - prob) ** gamma) * cross_entropy_loss_pos focal_pos *= alpha cross_entropy_loss_neg = criterion(inputs, torch.zeros_like(inputs)) focal_neg = (prob**gamma) * cross_entropy_loss_neg focal_neg *= 1 - alpha loss = torch.matmul(focal_pos, labels.T) + torch.matmul(focal_neg, (1 - labels).T) return loss / height_and_width class MaskFormerDetrDecoder(DetrDecoder): pass # refactored from original implementation class MaskFormerHungarianMatcher(nn.Module): """This class computes an assignment between the labels and the predictions of the network. For efficiency reasons, the labels don't include the no_object. Because of this, in general, there are more predictions than labels. In this case, we do a 1-to-1 matching of the best predictions, while the others are un-matched (and thus treated as non-objects). """ def __init__(self, cost_class: float = 1.0, cost_mask: float = 1.0, cost_dice: float = 1.0): """Creates the matcher Params: cost_class (float, *optional*, defaults to 1.0): This is the relative weight of the classification error in the matching cost. cost_mask (float, *optional*, defaults to 1.0): This is the relative weight of the focal loss of the binary mask in the matching cost. cost_dice (float, *optional*, defaults to 1.0): This is the relative weight of the dice loss of the binary mask in the matching cost """ super().__init__() if cost_class == 0 and cost_mask == 0 and cost_dice == 0: raise ValueError("All costs can't be 0") self.cost_class = cost_class self.cost_mask = cost_mask self.cost_dice = cost_dice @torch.no_grad() def forward(self, masks_queries_logits, class_queries_logits, mask_labels, class_labels) -> list[tuple[Tensor]]: """Performs the matching Invalid predictions or targets resulting in NaN or inf values in the matcher cost matrix do not raise errors and instead will only be assigned if no other valid prediction or target can be matched instead. This avoids random crashes at training time. A high training loss indicates that some of the predictions or targets might be invalid. If the loss doesn't improve after a couple of steps the model has likely diverged. Params: masks_queries_logits (`torch.Tensor`): A tensor` of dim `batch_size, num_queries, num_labels` with the classification logits. class_queries_logits (`torch.Tensor`): A tensor` of dim `batch_size, num_queries, height, width` with the predicted masks. class_labels (`torch.Tensor`): A tensor` of dim `num_target_boxes` (where num_target_boxes is the number of ground-truth objects in the target) containing the class labels. mask_labels (`torch.Tensor`): A tensor` of dim `num_target_boxes, height, width` containing the target masks. Returns: `list[tuple[Tensor]]`: A list of size batch_size, containing tuples of (index_i, index_j) where: - index_i is the indices of the selected predictions (in order) - index_j is the indices of the corresponding selected labels (in order) For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes). """ indices: list[tuple[np.array]] = [] preds_masks = masks_queries_logits preds_probs = class_queries_logits # iterate through batch size for pred_probs, pred_mask, target_mask, labels in zip(preds_probs, preds_masks, mask_labels, class_labels): # downsample the target mask, save memory target_mask = nn.functional.interpolate(target_mask[:, None], size=pred_mask.shape[-2:], mode="nearest") pred_probs = pred_probs.softmax(-1) # Compute the classification cost. Contrary to the loss, we don't use the NLL, # but approximate it in 1 - proba[target class]. # The 1 is a constant that doesn't change the matching, it can be omitted. cost_class = -pred_probs[:, labels] # flatten spatial dimension "q h w -> q (h w)" pred_mask_flat = pred_mask.flatten(1) # [num_queries, height*width] # same for target_mask "c h w -> c (h w)" target_mask_flat = target_mask[:, 0].flatten(1) # [num_total_labels, height*width] # compute the focal loss between each mask pairs -> shape (num_queries, num_labels) cost_mask = pair_wise_sigmoid_focal_loss(pred_mask_flat, target_mask_flat) # Compute the dice loss between each mask pairs -> shape (num_queries, num_labels) cost_dice = pair_wise_dice_loss(pred_mask_flat, target_mask_flat) # final cost matrix cost_matrix = self.cost_mask * cost_mask + self.cost_class * cost_class + self.cost_dice * cost_dice # Replace NaN and inf values with max value to avoid linear_sum_assignment errors. Max value is used to match # these predictions only if there are no other valid predictions. max_value = torch.finfo(cost_matrix.dtype).max cost_matrix = torch.nan_to_num(cost_matrix, nan=max_value, posinf=max_value, neginf=max_value) # do the assignment using the hungarian algorithm in scipy assigned_indices: tuple[np.array] = linear_sum_assignment(cost_matrix.cpu()) indices.append(assigned_indices) # It could be stacked in one tensor matched_indices = [ (torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices ] return matched_indices def __repr__(self): head = "Matcher " + self.__class__.__name__ body = [ f"cost_class: {self.cost_class}", f"cost_mask: {self.cost_mask}", f"cost_dice: {self.cost_dice}", ] _repr_indent = 4 lines = [head] + [" " * _repr_indent + line for line in body] return "\n".join(lines) # copied and adapted from original implementation class MaskFormerLoss(nn.Module): def __init__( self, num_labels: int, matcher: MaskFormerHungarianMatcher, weight_dict: dict[str, float], eos_coef: float, ): """ The MaskFormer Loss. The loss is computed very similar to DETR. The process happens in two steps: 1) we compute hungarian assignment between ground truth masks and the outputs of the model 2) we supervise each pair of matched ground-truth / prediction (supervise class and mask) Args: num_labels (`int`): The number of classes. matcher (`MaskFormerHungarianMatcher`): A torch module that computes the assignments between the predictions and labels. weight_dict (`dict[str, float]`): A dictionary of weights to be applied to the different losses. eos_coef (`float`): Weight to apply to the null class. """ super().__init__() requires_backends(self, ["scipy"]) self.num_labels = num_labels self.matcher = matcher self.weight_dict = weight_dict self.eos_coef = eos_coef empty_weight = torch.ones(self.num_labels + 1) empty_weight[-1] = self.eos_coef self.empty_weight = nn.Buffer(empty_weight) def _max_by_axis(self, the_list: list[list[int]]) -> list[int]: maxes = the_list[0] for sublist in the_list[1:]: for index, item in enumerate(sublist): maxes[index] = max(maxes[index], item) return maxes def _pad_images_to_max_in_batch(self, tensors: list[Tensor]) -> tuple[Tensor, Tensor]: # get the maximum size in the batch max_size = self._max_by_axis([list(tensor.shape) for tensor in tensors]) batch_size = len(tensors) # compute finel size batch_shape = [batch_size] + max_size b, _, h, w = batch_shape # get metadata dtype = tensors[0].dtype device = tensors[0].device padded_tensors = torch.zeros(batch_shape, dtype=dtype, device=device) padding_masks = torch.ones((b, h, w), dtype=torch.bool, device=device) # pad the tensors to the size of the biggest one for tensor, padded_tensor, padding_mask in zip(tensors, padded_tensors, padding_masks): padded_tensor[: tensor.shape[0], : tensor.shape[1], : tensor.shape[2]].copy_(tensor) padding_mask[: tensor.shape[1], : tensor.shape[2]] = False return padded_tensors, padding_masks def loss_labels( self, class_queries_logits: Tensor, class_labels: list[Tensor], indices: tuple[np.array] ) -> dict[str, Tensor]: """Compute the losses related to the labels using cross entropy. Args: class_queries_logits (`torch.Tensor`): A tensor of shape `batch_size, num_queries, num_labels` class_labels (`list[torch.Tensor]`): List of class labels of shape `(labels)`. indices (`tuple[np.array])`: The indices computed by the Hungarian matcher. Returns: `dict[str, Tensor]`: A dict of `torch.Tensor` containing the following key: - **loss_cross_entropy** -- The loss computed using cross entropy on the predicted and ground truth labels. """ pred_logits = class_queries_logits batch_size, num_queries, _ = pred_logits.shape criterion = nn.CrossEntropyLoss(weight=self.empty_weight) idx = self._get_predictions_permutation_indices(indices) # shape = (batch_size, num_queries) target_classes_o = torch.cat([target[j] for target, (_, j) in zip(class_labels, indices)]) # shape = (batch_size, num_queries) target_classes = torch.full( (batch_size, num_queries), fill_value=self.num_labels, dtype=torch.int64, device=pred_logits.device ) target_classes[idx] = target_classes_o # target_classes is a (batch_size, num_labels, num_queries), we need to permute pred_logits "b q c -> b c q" pred_logits_transposed = pred_logits.transpose(1, 2) loss_ce = criterion(pred_logits_transposed, target_classes) losses = {"loss_cross_entropy": loss_ce} return losses def loss_masks( self, masks_queries_logits: Tensor, mask_labels: list[Tensor], indices: tuple[np.array], num_masks: int ) -> dict[str, Tensor]: """Compute the losses related to the masks using focal and dice loss. Args: masks_queries_logits (`torch.Tensor`): A tensor of shape `batch_size, num_queries, height, width` mask_labels (`torch.Tensor`): List of mask labels of shape `(labels, height, width)`. indices (`tuple[np.array])`: The indices computed by the Hungarian matcher. num_masks (`int)`: The number of masks, used for normalization. Returns: `dict[str, Tensor]`: A dict of `torch.Tensor` containing two keys: - **loss_mask** -- The loss computed using sigmoid focal loss on the predicted and ground truth masks. - **loss_dice** -- The loss computed using dice loss on the predicted on the predicted and ground truth masks. """ src_idx = self._get_predictions_permutation_indices(indices) tgt_idx = self._get_targets_permutation_indices(indices) # shape (batch_size * num_queries, height, width) pred_masks = masks_queries_logits[src_idx] # shape (batch_size, num_queries, height, width) # pad all and stack the targets to the num_labels dimension target_masks, _ = self._pad_images_to_max_in_batch(mask_labels) target_masks = target_masks[tgt_idx] # upsample predictions to the target size, we have to add one dim to use interpolate pred_masks = nn.functional.interpolate( pred_masks[:, None], size=target_masks.shape[-2:], mode="bilinear", align_corners=False ) pred_masks = pred_masks[:, 0].flatten(1) target_masks = target_masks.flatten(1) losses = { "loss_mask": sigmoid_focal_loss(pred_masks, target_masks, num_masks), "loss_dice": dice_loss(pred_masks, target_masks, num_masks), } return losses def _get_predictions_permutation_indices(self, indices): # permute predictions following indices batch_indices = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)]) predictions_indices = torch.cat([src for (src, _) in indices]) return batch_indices, predictions_indices def _get_targets_permutation_indices(self, indices): # permute labels following indices batch_indices = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)]) target_indices = torch.cat([tgt for (_, tgt) in indices]) return batch_indices, target_indices def forward( self, masks_queries_logits: Tensor, class_queries_logits: Tensor, mask_labels: list[Tensor], class_labels: list[Tensor], auxiliary_predictions: dict[str, Tensor] | None = None, ) -> dict[str, Tensor]: """ This performs the loss computation. Args: masks_queries_logits (`torch.Tensor`): A tensor of shape `batch_size, num_queries, height, width` class_queries_logits (`torch.Tensor`): A tensor of shape `batch_size, num_queries, num_labels` mask_labels (`torch.Tensor`): List of mask labels of shape `(labels, height, width)`. class_labels (`list[torch.Tensor]`): List of class labels of shape `(labels)`. auxiliary_predictions (`dict[str, torch.Tensor]`, *optional*): if `use_auxiliary_loss` was set to `true` in [`MaskFormerConfig`], then it contains the logits from the inner layers of the Detr's Decoder. Returns: `dict[str, Tensor]`: A dict of `torch.Tensor` containing two keys: - **loss_cross_entropy** -- The loss computed using cross entropy on the predicted and ground truth labels. - **loss_mask** -- The loss computed using sigmoid focal loss on the predicted and ground truth masks. - **loss_dice** -- The loss computed using dice loss on the predicted on the predicted and ground truth masks. if `use_auxiliary_loss` was set to `true` in [`MaskFormerConfig`], the dictionary contains additional losses for each auxiliary predictions. """ # retrieve the matching between the outputs of the last layer and the labels indices = self.matcher(masks_queries_logits, class_queries_logits, mask_labels, class_labels) # compute the average number of target masks for normalization purposes num_masks: Number = self.get_num_masks(class_labels, device=class_labels[0].device) # get all the losses losses: dict[str, Tensor] = { **self.loss_masks(masks_queries_logits, mask_labels, indices, num_masks), **self.loss_labels(class_queries_logits, class_labels, indices), } # in case of auxiliary losses, we repeat this process with the output of each intermediate layer. if auxiliary_predictions is not None: for idx, aux_outputs in enumerate(auxiliary_predictions): masks_queries_logits = aux_outputs["masks_queries_logits"] class_queries_logits = aux_outputs["class_queries_logits"] loss_dict = self.forward(masks_queries_logits, class_queries_logits, mask_labels, class_labels) loss_dict = {f"{key}_{idx}": value for key, value in loss_dict.items()} losses.update(loss_dict) return losses def get_num_masks(self, class_labels: torch.Tensor, device: torch.device) -> torch.Tensor: """ Computes the average number of target masks across the batch, for normalization purposes. """ num_masks = sum(len(classes) for classes in class_labels) num_masks = torch.as_tensor(num_masks, dtype=torch.float, device=device) world_size = 1 if is_accelerate_available(): if PartialState._shared_state != {}: num_masks = reduce(num_masks) world_size = PartialState().num_processes num_masks = torch.clamp(num_masks / world_size, min=1) return num_masks class MaskFormerFPNConvLayer(nn.Module): def __init__(self, in_features: int, out_features: int, kernel_size: int = 3, padding: int = 1): """ A basic module that executes conv - norm - in sequence used in MaskFormer. Args: in_features (`int`): The number of input features (channels). out_features (`int`): The number of outputs features (channels). """ super().__init__() self.layers = [ nn.Conv2d(in_features, out_features, kernel_size=kernel_size, padding=padding, bias=False), nn.GroupNorm(32, out_features), nn.ReLU(inplace=True), ] for i, layer in enumerate(self.layers): # Provide backwards compatibility from when the class inherited from nn.Sequential # In nn.Sequential subclasses, the name given to the layer is its index in the sequence. # In nn.Module subclasses they derived from the instance attribute they are assigned to e.g. # self.my_layer_name = Layer() # We can't give instance attributes integer names i.e. self.0 is not permitted and so need to register # explicitly self.add_module(str(i), layer) def forward(self, input: Tensor) -> Tensor: hidden_state = input for layer in self.layers: hidden_state = layer(hidden_state) return hidden_state class MaskFormerFPNLayer(nn.Module): def __init__(self, in_features: int, lateral_features: int): """ A Feature Pyramid Network Layer (FPN) layer. It creates a feature map by aggregating features from the previous and backbone layer. Due to the spatial mismatch, the tensor coming from the previous layer is upsampled. Args: in_features (`int`): The number of input features (channels). lateral_features (`int`): The number of lateral features (channels). """ super().__init__() self.proj = nn.Sequential( nn.Conv2d(lateral_features, in_features, kernel_size=1, padding=0, bias=False), nn.GroupNorm(32, in_features), ) self.block = MaskFormerFPNConvLayer(in_features, in_features) def forward(self, down: Tensor, left: Tensor) -> Tensor: left = self.proj(left) down = nn.functional.interpolate(down, size=left.shape[-2:], mode="nearest") down += left down = self.block(down) return down class MaskFormerFPNModel(nn.Module): def __init__(self, in_features: int, lateral_widths: list[int], feature_size: int = 256): """ Feature Pyramid Network, given an input tensor and a set of feature map of different feature/spatial size, it creates a list of feature maps with the same feature size. Args: in_features (`int`): The number of input features (channels). lateral_widths (`list[int]`): A list with the features (channels) size of each lateral connection. feature_size (int, *optional*, defaults to 256): The features (channels) of the resulting feature maps. """ super().__init__() self.stem = MaskFormerFPNConvLayer(in_features, feature_size) self.layers = nn.Sequential( *[MaskFormerFPNLayer(feature_size, lateral_width) for lateral_width in lateral_widths[::-1]] ) def forward(self, features: list[Tensor]) -> list[Tensor]: fpn_features = [] last_feature = features[-1] other_features = features[:-1] output = self.stem(last_feature) for layer, left in zip(self.layers, other_features[::-1]): output = layer(output, left) fpn_features.append(output) return fpn_features class MaskFormerPixelDecoder(nn.Module): def __init__(self, *args, feature_size: int = 256, mask_feature_size: int = 256, **kwargs): r""" Pixel Decoder Module proposed in [Per-Pixel Classification is Not All You Need for Semantic Segmentation](https://huggingface.co/papers/2107.06278). It first runs the backbone's features into a Feature Pyramid Network creating a list of feature maps. Then, it projects the last one to the correct `mask_size`. Args: feature_size (`int`, *optional*, defaults to 256): The feature size (channel dimension) of the FPN feature maps. mask_feature_size (`int`, *optional*, defaults to 256): The features (channels) of the target masks size \\(C_{\epsilon}\\) in the paper. """ super().__init__() self.fpn = MaskFormerFPNModel(*args, feature_size=feature_size, **kwargs) self.mask_projection = nn.Conv2d(feature_size, mask_feature_size, kernel_size=3, padding=1) def forward( self, features: list[Tensor], output_hidden_states: bool = False, return_dict: bool = True ) -> MaskFormerPixelDecoderOutput: fpn_features = self.fpn(features) # we use the last feature map last_feature_projected = self.mask_projection(fpn_features[-1]) if not return_dict: return (last_feature_projected, tuple(fpn_features)) if output_hidden_states else (last_feature_projected,) return MaskFormerPixelDecoderOutput( last_hidden_state=last_feature_projected, hidden_states=tuple(fpn_features) if output_hidden_states else () ) class MaskFormerSinePositionEmbedding(DetrSinePositionEmbedding): pass class PredictionBlock(nn.Module): def __init__(self, in_dim: int, out_dim: int, activation: nn.Module) -> None: super().__init__() self.layers = [nn.Linear(in_dim, out_dim), activation] # Maintain submodule indexing as if part of a Sequential block for i, layer in enumerate(self.layers): self.add_module(str(i), layer) def forward(self, input: Tensor) -> Tensor: hidden_state = input for layer in self.layers: hidden_state = layer(hidden_state) return hidden_state class MaskformerMLPPredictionHead(nn.Module): def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int = 3): """ A classic Multi Layer Perceptron (MLP). Args: input_dim (`int`): The input dimensions. hidden_dim (`int`): The hidden dimensions. output_dim (`int`): The output dimensions. num_layers (int, *optional*, defaults to 3): The number of layers. """ super().__init__() in_dims = [input_dim] + [hidden_dim] * (num_layers - 1) out_dims = [hidden_dim] * (num_layers - 1) + [output_dim] self.layers = [] for i, (in_dim, out_dim) in enumerate(zip(in_dims, out_dims)): activation = nn.ReLU() if i < num_layers - 1 else nn.Identity() layer = PredictionBlock(in_dim, out_dim, activation=activation) self.layers.append(layer) # Provide backwards compatibility from when the class inherited from nn.Sequential # In nn.Sequential subclasses, the name given to the layer is its index in the sequence. # In nn.Module subclasses they derived from the instance attribute they are assigned to e.g. # self.my_layer_name = Layer() # We can't give instance attributes integer names i.e. self.0 is not permitted and so need to register # explicitly self.add_module(str(i), layer) def forward(self, input: Tensor) -> Tensor: hidden_state = input for layer in self.layers: hidden_state = layer(hidden_state) return hidden_state class MaskFormerPixelLevelModule(nn.Module): def __init__(self, config: MaskFormerConfig): """ Pixel Level Module proposed in [Per-Pixel Classification is Not All You Need for Semantic Segmentation](https://huggingface.co/papers/2107.06278). It runs the input image through a backbone and a pixel decoder, generating an image feature map and pixel embeddings. Args: config ([`MaskFormerConfig`]): The configuration used to instantiate this model. """ super().__init__() if getattr(config, "backbone_config") is not None and config.backbone_config.model_type == "swin": # for backwards compatibility backbone_config = config.backbone_config backbone_config = MaskFormerSwinConfig.from_dict(backbone_config.to_dict()) backbone_config.out_features = ["stage1", "stage2", "stage3", "stage4"] config.backbone_config = backbone_config self.encoder = AutoBackbone.from_config(config=config.backbone_config) feature_channels = self.encoder.channels self.decoder = MaskFormerPixelDecoder( in_features=feature_channels[-1], feature_size=config.fpn_feature_size, mask_feature_size=config.mask_feature_size, lateral_widths=feature_channels[:-1], ) def forward( self, pixel_values: Tensor, output_hidden_states: bool = False, return_dict: bool = True ) -> MaskFormerPixelLevelModuleOutput: features = self.encoder(pixel_values).feature_maps decoder_output = self.decoder(features, output_hidden_states, return_dict=return_dict) if not return_dict: last_hidden_state = decoder_output[0] outputs = (features[-1], last_hidden_state) if output_hidden_states: hidden_states = decoder_output[1] outputs = outputs + (tuple(features),) + (hidden_states,) return outputs return MaskFormerPixelLevelModuleOutput( # the last feature is actually the output from the last layer encoder_last_hidden_state=features[-1], decoder_last_hidden_state=decoder_output.last_hidden_state, encoder_hidden_states=tuple(features) if output_hidden_states else (), decoder_hidden_states=decoder_output.hidden_states if output_hidden_states else (), ) class MaskFormerTransformerModule(nn.Module): """ The MaskFormer's transformer module. """ def __init__(self, in_features: int, config: MaskFormerConfig): super().__init__() hidden_size = config.decoder_config.hidden_size should_project = in_features != hidden_size self.position_embedder = MaskFormerSinePositionEmbedding( num_position_features=hidden_size // 2, normalize=True ) self.queries_embedder = nn.Embedding(config.decoder_config.num_queries, hidden_size) self.input_projection = nn.Conv2d(in_features, hidden_size, kernel_size=1) if should_project else None self.decoder = MaskFormerDetrDecoder(config=config.decoder_config) def forward( self, image_features: Tensor, output_hidden_states: bool = False, output_attentions: bool = False, return_dict: bool | None = None, ) -> DetrDecoderOutput: if self.input_projection is not None: image_features = self.input_projection(image_features) # repeat the queries "q c -> b q c" batch_size = image_features.shape[0] queries_embeddings = self.queries_embedder.weight.unsqueeze(0).repeat(batch_size, 1, 1) inputs_embeds = torch.zeros_like(queries_embeddings, requires_grad=self.training) # torch.export.export does no support requires_grad if self.training: inputs_embeds.requires_grad_(True) batch_size, num_channels, height, width = image_features.shape # rearrange both image_features and object_queries "b c h w -> b (h w) c" object_queries = ( self.position_embedder(image_features.shape, image_features.device, image_features.dtype) .view(batch_size, num_channels, height * width) .transpose(1, 2) ) image_features = image_features.view(batch_size, num_channels, height * width).transpose(1, 2) decoder_output: DetrDecoderOutput = self.decoder( inputs_embeds=inputs_embeds, attention_mask=None, encoder_hidden_states=image_features, encoder_attention_mask=None, spatial_position_embeddings=object_queries, object_queries_position_embeddings=queries_embeddings, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) return decoder_output @auto_docstring class MaskFormerPreTrainedModel(PreTrainedModel): config: MaskFormerConfig base_model_prefix = "model" main_input_name = "pixel_values" input_modalities = ("image",) @torch.no_grad() def _init_weights(self, module: nn.Module): super()._init_weights(module) xavier_std = self.config.init_xavier_std if isinstance(module, MaskFormerTransformerModule): if module.input_projection is not None: init.xavier_uniform_(module.input_projection.weight, gain=xavier_std) init.constant_(module.input_projection.bias, 0) # FPN elif isinstance(module, MaskFormerFPNModel): init.xavier_uniform_(module.stem.get_submodule("0").weight, gain=xavier_std) elif isinstance(module, MaskFormerFPNLayer): init.xavier_uniform_(module.proj[0].weight, gain=xavier_std) elif isinstance(module, MaskFormerFPNConvLayer): init.xavier_uniform_(module.get_submodule("0").weight, gain=xavier_std) # The MLP head elif isinstance(module, MaskformerMLPPredictionHead): # I was not able to find the correct initializer in the original implementation # we'll use xavier for submodule in module.modules(): if isinstance(submodule, nn.Linear): init.xavier_uniform_(submodule.weight, gain=xavier_std) init.constant_(submodule.bias, 0) # The base class treats BatchNorm as a norm (weight=ones_); MaskFormer instead applies a # normal_ init to the weight and resets the running stats, so this branch must be kept. elif isinstance(module, nn.BatchNorm2d): init.normal_(module.weight, mean=0.0, std=self.config.init_std) if module.bias is not None: init.zeros_(module.bias) if getattr(module, "running_mean", None) is not None: init.zeros_(module.running_mean) init.ones_(module.running_var) init.zeros_(module.num_batches_tracked) elif isinstance(module, MaskFormerLoss): empty_weight = torch.ones(module.num_labels + 1) empty_weight[-1] = module.eos_coef init.copy_(module.empty_weight, empty_weight) @auto_docstring class MaskFormerModel(MaskFormerPreTrainedModel): def __init__(self, config: MaskFormerConfig): super().__init__(config) self.pixel_level_module = MaskFormerPixelLevelModule(config) self.transformer_module = MaskFormerTransformerModule( in_features=self.pixel_level_module.encoder.channels[-1], config=config ) self.post_init() @auto_docstring def forward( self, pixel_values: Tensor, pixel_mask: Tensor | None = None, output_hidden_states: bool | None = None, output_attentions: bool | None = None, return_dict: bool | None = None, **kwargs, ) -> MaskFormerModelOutput: r""" Examples: ```python >>> from transformers import AutoImageProcessor, MaskFormerModel >>> from PIL import Image >>> from huggingface_hub.utils import httpx >>> from io import BytesIO >>> # load MaskFormer fine-tuned on ADE20k semantic segmentation >>> image_processor = AutoImageProcessor.from_pretrained("facebook/maskformer-swin-base-ade") >>> model = MaskFormerModel.from_pretrained("facebook/maskformer-swin-base-ade") >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> with httpx.stream("GET", url) as response: ... image = Image.open(BytesIO(response.read())) >>> inputs = image_processor(image, return_tensors="pt") >>> # forward pass >>> outputs = model(**inputs) >>> # the decoder of MaskFormer outputs hidden states of shape (batch_size, num_queries, hidden_size) >>> transformer_decoder_last_hidden_state = outputs.transformer_decoder_last_hidden_state >>> list(transformer_decoder_last_hidden_state.shape) [1, 100, 256] ```""" if pixel_values is None: raise ValueError("You have to specify pixel_values") output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.return_dict batch_size, _, height, width = pixel_values.shape if pixel_mask is None: pixel_mask = torch.ones((batch_size, height, width), device=pixel_values.device) pixel_level_module_output = self.pixel_level_module( pixel_values, output_hidden_states, return_dict=return_dict ) image_features = pixel_level_module_output[0] pixel_embeddings = pixel_level_module_output[1] transformer_module_output = self.transformer_module(image_features, output_hidden_states, output_attentions) queries = transformer_module_output.last_hidden_state encoder_hidden_states = None pixel_decoder_hidden_states = None transformer_decoder_hidden_states = None hidden_states = None if output_hidden_states: encoder_hidden_states = pixel_level_module_output[2] pixel_decoder_hidden_states = pixel_level_module_output[3] transformer_decoder_hidden_states = transformer_module_output[1] hidden_states = encoder_hidden_states + pixel_decoder_hidden_states + transformer_decoder_hidden_states output = MaskFormerModelOutput( encoder_last_hidden_state=image_features, pixel_decoder_last_hidden_state=pixel_embeddings, transformer_decoder_last_hidden_state=queries, encoder_hidden_states=encoder_hidden_states, pixel_decoder_hidden_states=pixel_decoder_hidden_states, transformer_decoder_hidden_states=transformer_decoder_hidden_states, hidden_states=hidden_states, attentions=transformer_module_output.attentions, ) if not return_dict: output = tuple(v for v in output.values()) return output class MaskFormerForInstanceSegmentation(MaskFormerPreTrainedModel): def __init__(self, config: MaskFormerConfig): super().__init__(config) self.model = MaskFormerModel(config) hidden_size = config.decoder_config.hidden_size # + 1 because we add the "null" class self.class_predictor = nn.Linear(hidden_size, config.num_labels + 1) self.mask_embedder = MaskformerMLPPredictionHead(hidden_size, hidden_size, config.mask_feature_size) self.matcher = MaskFormerHungarianMatcher( cost_class=1.0, cost_dice=config.dice_weight, cost_mask=config.mask_weight ) self.weight_dict: dict[str, float] = { "loss_cross_entropy": config.cross_entropy_weight, "loss_mask": config.mask_weight, "loss_dice": config.dice_weight, } self.criterion = MaskFormerLoss( config.num_labels, matcher=self.matcher, weight_dict=self.weight_dict, eos_coef=config.no_object_weight, ) self.post_init() def get_loss_dict( self, masks_queries_logits: Tensor, class_queries_logits: Tensor, mask_labels: Tensor, class_labels: Tensor, auxiliary_logits: dict[str, Tensor], ) -> dict[str, Tensor]: loss_dict: dict[str, Tensor] = self.criterion( masks_queries_logits, class_queries_logits, mask_labels, class_labels, auxiliary_logits ) # weight each loss by `self.weight_dict[]` including auxiliary losses for key, weight in self.weight_dict.items(): for loss_key, loss in loss_dict.items(): if key in loss_key: loss *= weight return loss_dict def get_loss(self, loss_dict: dict[str, Tensor]) -> Tensor: return sum(loss_dict.values()) def get_logits(self, outputs: MaskFormerModelOutput) -> tuple[Tensor, Tensor, dict[str, Tensor]]: pixel_embeddings = outputs.pixel_decoder_last_hidden_state # get the auxiliary predictions (one for each decoder's layer) auxiliary_logits: list[str, Tensor] = [] # This code is a little bit cumbersome, an improvement can be to return a list of predictions. If we have auxiliary loss then we are going to return more than one element in the list if self.config.use_auxiliary_loss: stacked_transformer_decoder_outputs = torch.stack(outputs.transformer_decoder_hidden_states) classes = self.class_predictor(stacked_transformer_decoder_outputs) class_queries_logits = classes[-1] # get the masks mask_embeddings = self.mask_embedder(stacked_transformer_decoder_outputs) binaries_masks = torch.einsum("lbqc, bchw -> lbqhw", mask_embeddings, pixel_embeddings) masks_queries_logits = binaries_masks[-1] # go til [:-1] because the last one is always used for aux_binary_masks, aux_classes in zip(binaries_masks[:-1], classes[:-1]): auxiliary_logits.append( {"masks_queries_logits": aux_binary_masks, "class_queries_logits": aux_classes} ) else: transformer_decoder_hidden_states = outputs.transformer_decoder_last_hidden_state classes = self.class_predictor(transformer_decoder_hidden_states) class_queries_logits = classes # get the masks mask_embeddings = self.mask_embedder(transformer_decoder_hidden_states) # sum up over the channels masks_queries_logits = torch.einsum("bqc, bchw -> bqhw", mask_embeddings, pixel_embeddings) return class_queries_logits, masks_queries_logits, auxiliary_logits @auto_docstring def forward( self, pixel_values: Tensor, mask_labels: list[Tensor] | None = None, class_labels: list[Tensor] | None = None, pixel_mask: Tensor | None = None, output_auxiliary_logits: bool | None = None, output_hidden_states: bool | None = None, output_attentions: bool | None = None, return_dict: bool | None = None, **kwargs, ) -> MaskFormerForInstanceSegmentationOutput: r""" mask_labels (`list[torch.Tensor]`, *optional*): List of mask labels of shape `(num_labels, height, width)` to be fed to a model class_labels (`list[torch.LongTensor]`, *optional*): list of target class labels of shape `(num_labels, height, width)` to be fed to a model. They identify the labels of `mask_labels`, e.g. the label of `mask_labels[i][j]` if `class_labels[i][j]`. output_auxiliary_logits (`bool`, *optional*): Whether or not to output auxiliary logits. Examples: Semantic segmentation example: ```python >>> from transformers import AutoImageProcessor, MaskFormerForInstanceSegmentation >>> from PIL import Image >>> from huggingface_hub.utils import httpx >>> from io import BytesIO >>> # load MaskFormer fine-tuned on ADE20k semantic segmentation >>> image_processor = AutoImageProcessor.from_pretrained("facebook/maskformer-swin-base-ade") >>> model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-base-ade") >>> url = ( ... "https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg" ... ) >>> with httpx.stream("GET", url) as response: ... image = Image.open(BytesIO(response.read())) >>> inputs = image_processor(images=image, return_tensors="pt") >>> outputs = model(**inputs) >>> # model predicts class_queries_logits of shape `(batch_size, num_queries)` >>> # and masks_queries_logits of shape `(batch_size, num_queries, height, width)` >>> class_queries_logits = outputs.class_queries_logits >>> masks_queries_logits = outputs.masks_queries_logits >>> # you can pass them to image_processor for postprocessing >>> predicted_semantic_map = image_processor.post_process_semantic_segmentation( ... outputs, target_sizes=[(image.height, image.width)] ... )[0] >>> # we refer to the demo notebooks for visualization (see "Resources" section in the MaskFormer docs) >>> list(predicted_semantic_map.shape) [512, 683] ``` Panoptic segmentation example: ```python >>> from transformers import AutoImageProcessor, MaskFormerForInstanceSegmentation >>> from PIL import Image >>> from huggingface_hub.utils import httpx >>> from io import BytesIO >>> # load MaskFormer fine-tuned on COCO panoptic segmentation >>> image_processor = AutoImageProcessor.from_pretrained("facebook/maskformer-swin-base-coco") >>> model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-base-coco") >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> with httpx.stream("GET", url) as response: ... image = Image.open(BytesIO(response.read())) >>> inputs = image_processor(images=image, return_tensors="pt") >>> outputs = model(**inputs) >>> # model predicts class_queries_logits of shape `(batch_size, num_queries)` >>> # and masks_queries_logits of shape `(batch_size, num_queries, height, width)` >>> class_queries_logits = outputs.class_queries_logits >>> masks_queries_logits = outputs.masks_queries_logits >>> # you can pass them to image_processor for postprocessing >>> result = image_processor.post_process_panoptic_segmentation(outputs, target_sizes=[(image.height, image.width)])[0] >>> # we refer to the demo notebooks for visualization (see "Resources" section in the MaskFormer docs) >>> predicted_panoptic_map = result["segmentation"] >>> list(predicted_panoptic_map.shape) [480, 640] ``` """ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.return_dict raw_outputs = self.model( pixel_values, pixel_mask, output_hidden_states=output_hidden_states or self.config.use_auxiliary_loss, return_dict=return_dict, output_attentions=output_attentions, ) # We need to have raw_outputs optionally be returned as a dict to use torch.compile. For backwards # compatibility we convert to a dataclass for the rest of the model logic outputs = MaskFormerModelOutput( encoder_last_hidden_state=raw_outputs[0], pixel_decoder_last_hidden_state=raw_outputs[1], transformer_decoder_last_hidden_state=raw_outputs[2], encoder_hidden_states=raw_outputs[3] if output_hidden_states else None, pixel_decoder_hidden_states=raw_outputs[4] if output_hidden_states else None, transformer_decoder_hidden_states=raw_outputs[5] if output_hidden_states else None, hidden_states=raw_outputs[6] if output_hidden_states else None, attentions=raw_outputs[-1] if output_attentions else None, ) loss, loss_dict, auxiliary_logits = None, None, None class_queries_logits, masks_queries_logits, auxiliary_logits = self.get_logits(outputs) if mask_labels is not None and class_labels is not None: loss_dict: dict[str, Tensor] = self.get_loss_dict( masks_queries_logits, class_queries_logits, mask_labels, class_labels, auxiliary_logits ) loss = self.get_loss(loss_dict) output_auxiliary_logits = ( self.config.output_auxiliary_logits if output_auxiliary_logits is None else output_auxiliary_logits ) if not output_auxiliary_logits: auxiliary_logits = None if not return_dict: output = tuple( v for v in (loss, class_queries_logits, masks_queries_logits, auxiliary_logits, *outputs.values()) if v is not None ) return output return MaskFormerForInstanceSegmentationOutput( loss=loss, **outputs, class_queries_logits=class_queries_logits, masks_queries_logits=masks_queries_logits, auxiliary_logits=auxiliary_logits, ) __all__ = [ "MaskFormerConfig", "MaskFormerDetrConfig", "MaskFormerForInstanceSegmentation", "MaskFormerModel", "MaskFormerPreTrainedModel", "MaskFormerDetrPreTrainedModel", # noqa F821 ]