code: # Copyright 2021 The TensorFlow Authors. 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.
  
  """Build simclr models."""
  from typing import List, Optional, Mapping
  from absl import logging
  
  import tensorflow as tf
  
  from official.vision.beta.ops import anchor
  
  layers = tf.keras.layers
  
  
  @tf.keras.utils.register_keras_serializable(package='yolact')
  class YolactModel(tf.keras.Model):
    """A classification model based on SimCLR framework."""
  
    def __init__(self,
                 backbone: tf.keras.models.Model,
                 decoder: tf.keras.models.Model,
                 prediction_head: tf.keras.layers.Layer,
                 protonet: tf.keras.layers.Layer,
                 detection_generator: tf.keras.layers.Layer,
                 min_level: int,
                 max_level: int,
                 num_scales: int,
                 aspect_ratios: List[float],
                 anchor_size: float,
                 **kwargs):
      """A classification model based on SimCLR framework.
      Args:
        backbone: a backbone network.
        input_specs: `tf.keras.layers.InputSpec` specs of the input tensor.
        mode: `str` indicates mode of training to be executed.
        backbone_trainable: `bool` whether the backbone is trainable or not.
        **kwargs: keyword arguments to be passed.
      """
      super(YolactModel, self).__init__(**kwargs)
      self._config_dict = {
          'backbone': backbone,
          'decoder': decoder,
          'prediction_head': prediction_head,
          'protonet': protonet,
          'detection_generator': detection_generator,
          'min_level': min_level,
          'max_level': max_level,
          'num_scales': num_scales,
          'aspect_ratios': aspect_ratios,
          'anchor_size': anchor_size,
      }
      self.backbone = backbone
      self.decoder = decoder
      self.prediction_head = prediction_head
      self.protonet = protonet
      self.detection_generator = detection_generator
  
    def call(self,
             inputs: tf.Tensor,
             image_shape: Optional[tf.Tensor] = None,
             anchor_boxes: Optional[Mapping[str, tf.Tensor]] = None,
             training=None,
             **kwargs):
      """Forward pass of the YOLACT model.
  
      Args:
        images: `Tensor`, the input batched images, whose shape is
          [batch, height, width, 3].
        image_shape: `Tensor`, the actual shape of the input images, whose shape
          is [batch, 2] where the last dimension is [height, width]. Note that
          this is the actual image shape excluding paddings. For example, images
          in the batch may be resized into different shapes before padding to the
          fixed size.
        anchor_boxes: a dict of tensors which includes multilevel anchors.
          - key: `str`, the level of the multilevel predictions.
          - values: `Tensor`, the anchor coordinates of a particular feature
              level, whose shape is [height_l, width_l, num_anchors_per_location].
        training: `bool`, indicating whether it is in training mode.
  
      Returns:
        scores: a dict of tensors which includes scores of the predictions.
          - key: `str`, the level of the multilevel predictions.
          - values: `Tensor`, the box scores predicted from a particular feature
              level, whose shape is
              [batch, height_l, width_l, num_classes * num_anchors_per_location].
        boxes: a dict of tensors which includes coordinates of the predictions.
          - key: `str`, the level of the multilevel predictions.
          - values: `Tensor`, the box coordinates predicted from a particular
              feature level, whose shape is
              [batch, height_l, width_l, 4 * num_anchors_per_location].
        masks: a dict of tensors which includes mask coefficients of the predictions.
          - key: `str`, the level of the multilevel predictions.
          - values: `Tensor`, the mask coefficients predicted from a particular
              feature level, whose shape is
              [batch, height_l, width_l, k * num_anchors_per_location].
        protonet_features: `Tensor`, the protonet features, whose shape is
          [batch, height_2, width_2, k].
      """
      outputs = {}
      backbone_features = self.backbone(inputs)
      decoder_features = self.decoder(backbone_features)
      levels = sorted(decoder_features.keys())  # Ascending order
  
      raw_scores, raw_boxes, raw_masks = self.prediction_head(decoder_features)
      protonet_features = self.protonet(decoder_features[levels[0]])
  
      if training:
        outputs.update({
            'cls_outputs': raw_scores,
            'box_outputs': raw_boxes,
            'mask_outputs': raw_masks,
            'protonet_features': protonet_features,
        })
      else:
        # Generate anchor boxes for this batch if not provided.
        if anchor_boxes is None:
          _, image_height, image_width, _ = inputs.get_shape().as_list()
          anchor_boxes = anchor.Anchor(
              min_level=self._config_dict['min_level'],
              max_level=self._config_dict['max_level'],
              num_scales=self._config_dict['num_scales'],
              aspect_ratios=self._config_dict['aspect_ratios'],
              anchor_size=self._config_dict['anchor_size'],
              image_size=(image_height, image_width)).multilevel_boxes
          for l in anchor_boxes:
            anchor_boxes[l] = tf.tile(
                tf.expand_dims(anchor_boxes[l], axis=0),
                [tf.shape(inputs)[0], 1, 1, 1])
        # Post-processing.
        raw_attributes = {
            'raw_masks' : raw_masks,
        }
        final_results = self.detection_generator(raw_boxes, raw_scores,
                                                 anchor_boxes, image_shape,
                                                 raw_attributes)
  
        outputs.update({
            'cls_outputs': raw_scores,
            'box_outputs': raw_boxes,
            'mask_outputs': raw_masks,
            'protonet_features': protonet_features,
        })
  
        outputs.update({
            'detection_boxes': final_results['detection_boxes'],
            'detection_scores': final_results['detection_scores'],
            'detection_classes': final_results['detection_classes'],
            'num_detections': final_results['num_detections']
        })
  
        final_mask_coefficients = final_results['detection_attributes']['raw_masks']
        batch, proto_height, proto_width, proto_channel = protonet_features.get_shape().as_list()
        protonet_features = tf.reshape(protonet_features, [batch, -1, proto_channel]) # [batch, H*W, 32]
        assembled_masks = tf.matmul(
            a=protonet_features, # [batch, proto_height*proto_width, proto_channel]
            b=final_mask_coefficients, # [batch, max_num_instances, 32]
            transpose_b=True)
        assembled_masks = tf.reshape(assembled_masks, [batch, proto_height, proto_width, -1])
        assembled_masks = tf.transpose(assembled_masks, perm=[0,3,1,2])
        # [batch, max_num_instances, proto_height, proto_width]
  
        outputs['detection_masks'] = assembled_masks
  
      
      return outputs
  
    @property
    def checkpoint_items(self):
      """Returns a dictionary of items to be additionally checkpointed."""
      items = dict(backbone=self.backbone)
      items.update(decoder=self.decoder)
      return items
  
    def get_config(self):
      return self._config_dict
  
    @classmethod
    def from_config(cls, config, custom_objects=None):
      return cls(**config)
  
apis: tensorflow.shape tensorflow.reshape tensorflow.expand_dims tensorflow.matmul tensorflow.keras.utils.register_keras_serializable tensorflow.transpose