text: # Copyright 2016-2021 The Van Valen Lab at the California Institute of
  # Technology (Caltech), with support from the Paul Allen Family Foundation,
  # Google, & National Institutes of Health (NIH) under Grant U24CA224309-01.
  # All rights reserved.
  #
  # Licensed under a modified 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.github.com/vanvalenlab/tf-keras-retinanet/LICENSE
  #
  # The Work provided may be used for non-commercial academic purposes only.
  # For any other use of the Work, including commercial use, please contact:
  # vanvalenlab@gmail.com
  #
  # Neither the name of Caltech nor the names of its contributors may be used
  # to endorse or promote products derived from this software without specific
  # prior written permission.
  #
  # 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.
  # ==============================================================================
  """Upsampling layers"""
  from __future__ import absolute_import
  from __future__ import print_function
  from __future__ import division
  
  import tensorflow as tf
  from tensorflow.python.framework import tensor_shape
  from tensorflow.keras.layers import Layer
  from tensorflow.keras import backend as K
  from tensorflow.python.keras.utils import conv_utils
  
  
  class UpsampleLike(Layer):
      """Layer for upsampling a Tensor to be the same shape as another Tensor.
  
      Adapted from https://github.com/fizyr/keras-retinanet.
  
      Args:
          data_format (str): A string, one of ``channels_last`` (default)
              or ``channels_first``. The ordering of the dimensions in the
              inputs. ``channels_last`` corresponds to inputs with shape
              ``(batch, height, width, channels)`` while ``channels_first``
              corresponds to inputs with shape
              ``(batch, channels, height, width)``.
      """
  
      def __init__(self, data_format=None, **kwargs):
          super(UpsampleLike, self).__init__(**kwargs)
          self.data_format = conv_utils.normalize_data_format(data_format)
  
      def _resize_drop_axis(self, image, size, axis):
          image_shape = tf.shape(image)
  
          new_shape = []
          axes_resized = list(set([0, 1, 2, 3, 4]) - set([0, 4, axis]))
          for ax in range(K.ndim(image) - 1):
              if ax != axis:
                  new_shape.append(image_shape[ax])
              if ax == 3:
                  new_shape.append(image_shape[-1] * image_shape[axis])
  
          new_shape_2 = []
          for ax in range(K.ndim(image)):
              if ax == 0 or ax == 4 or ax == axis:
                  new_shape_2.append(image_shape[ax])
              elif ax == axes_resized[0]:
                  new_shape_2.append(size[0])
              elif ax == axes_resized[1]:
                  new_shape_2.append(size[1])
  
          new_image = tf.reshape(image, new_shape)
          new_image_resized = tf.image.resize(
              new_image,
              size,
              method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
  
          new_image_2 = tf.reshape(new_image_resized, new_shape_2)
  
          return new_image_2
  
      def resize_volumes(self, volume, size):
          # TODO: K.resize_volumes?
          if self.data_format == 'channels_first':
              volume = tf.transpose(volume, (0, 2, 3, 4, 1))
              new_size = (size[2], size[3], size[4])
          else:
              new_size = (size[1], size[2], size[3])
  
          new_shape_0 = (new_size[1], new_size[2])
          new_shape_1 = (new_size[0], new_size[1])
  
          resized_volume = self._resize_drop_axis(volume, new_shape_0, axis=1)
          resized_volume = self._resize_drop_axis(resized_volume, new_shape_1, axis=3)
  
          new_shape_static = [None, None, None, None, volume.get_shape()[-1]]
          resized_volume.set_shape(new_shape_static)
  
          if self.data_format == 'channels_first':
              resized_volume = tf.transpose(resized_volume, (0, 4, 1, 2, 3))
  
          return resized_volume
  
      def call(self, inputs, **kwargs):
          source, target = inputs
          target_shape = K.shape(target)
          if source.get_shape().ndims == 4:
              if self.data_format == 'channels_first':
                  source = tf.transpose(source, (0, 2, 3, 1))
                  new_shape = (target_shape[2], target_shape[3])
                  # TODO: K.resize_images?
                  output = tf.image.resize(
                      source, new_shape,
                      method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
                  output = tf.transpose(output, (0, 3, 1, 2))
                  return output
              new_shape = (target_shape[1], target_shape[2])
              return tf.image.resize(
                  source, new_shape,
                  method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
  
          if source.get_shape().ndims == 5:
              output = self.resize_volumes(source, target_shape)
              return output
  
          else:
              raise ValueError('Expected input[0] to have ndim of 4 or 5, found'
                               ' %s.' % source.get_shape().ndims)
  
      def compute_output_shape(self, input_shape):
          in_0 = tensor_shape.TensorShape(input_shape[0]).as_list()
          in_1 = tensor_shape.TensorShape(input_shape[1]).as_list()
          if self.data_format == 'channels_first':
              return tensor_shape.TensorShape(([in_0[0], in_0[1]] + in_1[2:]))
          return tensor_shape.TensorShape(([in_0[0]] + in_1[1:-1] + [in_0[-1]]))
  
      def get_config(self):
          config = {'data_format': self.data_format}
          base_config = super(UpsampleLike, self).get_config()
          return dict(list(base_config.items()) + list(config.items()))
  