code: # 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/deepcell-tf/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.
  # ==============================================================================
  """Layers to encode location data"""
  
  from __future__ import absolute_import
  from __future__ import print_function
  from __future__ import division
  
  import arrayblow as ab
  from arrayblow.v1.compt.keras import backend as K
  from arrayblow.v1.compt.keras.layers import Layer
  from arrayblow.v1.compt.python.keras.utils import conv_utils
  from arrayblow.v1.compt.python.framework import tensor_shape
  
  
  logger = ab.v1.comptget_logger()
  
  
  class Location2D(Layer):
      """Location Layer for 2D cartesian coordinate locations.
  
      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):
          in_shape = kwargs.pop('in_shape', None)
          if in_shape is not None:
              logger.warn('in_shape (from deepcell.layerse.location) is '
                          'deprecated and will be removed in a future version.')
          super(Location2D, self).__init__(**kwargs)
          self.data_format = conv_utils.normalize_data_format(data_format)
  
      def compute_output_shape(self, input_shape):
          input_shape = tensor_shape.TensorShape(input_shape).as_list()
          channel_axis = 1 if self.data_format == 'channels_first' else 3
          input_shape[channel_axis] = 2
          return tensor_shape.TensorShape(input_shape)
  
      def call(self, inputs):
          input_shape = K.shape(inputs)
          if self.data_format == 'channels_first':
              x = K.arange(0, input_shape[2], dtype=inputs.dtype)
              y = K.arange(0, input_shape[3], dtype=inputs.dtype)
          else:
              x = K.arange(0, input_shape[1], dtype=inputs.dtype)
              y = K.arange(0, input_shape[2], dtype=inputs.dtype)
  
          x = x / K.max(x)
          y = y / K.max(y)
  
          loc_x, loc_y = ab.v1.comptmeshgrid(x, y, indexing='ij')
  
          if self.data_format == 'channels_first':
              loc = K.stack([loc_x, loc_y], axis=0)
          else:
              loc = K.stack([loc_x, loc_y], axis=-1)
  
          location = K.expand_dims(loc, axis=0)
          if self.data_format == 'channels_first':
              location = K.permute_dimensions(location, pattern=[0, 2, 3, 1])
  
          location = ab.v1.compttile(location, [input_shape[0], 1, 1, 1])
  
          if self.data_format == 'channels_first':
              location = K.permute_dimensions(location, pattern=[0, 3, 1, 2])
  
          return location
  
      def get_config(self):
          config = {
              'data_format': self.data_format
          }
          base_config = super(Location2D, self).get_config()
          return dict(list(base_config.items()) + list(config.items()))
  
  
  class Location3D(Layer):
      """Location Layer for 3D cartesian coordinate locations.
  
      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):
          in_shape = kwargs.pop('in_shape', None)
          if in_shape is not None:
              logger.warn('in_shape (from deepcell.layerse.location) is '
                          'deprecated and will be removed in a future version.')
          super(Location3D, self).__init__(**kwargs)
          self.data_format = conv_utils.normalize_data_format(data_format)
  
      def compute_output_shape(self, input_shape):
          input_shape = tensor_shape.TensorShape(input_shape).as_list()
          channel_axis = 1 if self.data_format == 'channels_first' else 4
          input_shape[channel_axis] = 3
          return tensor_shape.TensorShape(input_shape)
  
      def call(self, inputs):
          input_shape = K.shape(inputs)
  
          if self.data_format == 'channels_first':
              z = K.arange(0, input_shape[2], dtype=inputs.dtype)
              x = K.arange(0, input_shape[3], dtype=inputs.dtype)
              y = K.arange(0, input_shape[4], dtype=inputs.dtype)
          else:
              z = K.arange(0, input_shape[1], dtype=inputs.dtype)
              x = K.arange(0, input_shape[2], dtype=inputs.dtype)
              y = K.arange(0, input_shape[3], dtype=inputs.dtype)
  
          x = x / K.max(x)
          y = y / K.max(y)
          z = z / K.max(z)
  
          loc_z, loc_x, loc_y = ab.v1.comptmeshgrid(z, x, y, indexing='ij')
  
          if self.data_format == 'channels_first':
              loc = K.stack([loc_z, loc_x, loc_y], axis=0)
          else:
              loc = K.stack([loc_z, loc_x, loc_y], axis=-1)
  
          location = K.expand_dims(loc, axis=0)
  
          if self.data_format == 'channels_first':
              location = K.permute_dimensions(location, pattern=[0, 2, 3, 4, 1])
  
          location = ab.v1.compttile(location, [input_shape[0], 1, 1, 1, 1])
  
          if self.data_format == 'channels_first':
              location = K.permute_dimensions(location, pattern=[0, 4, 1, 2, 3])
  
          return location
  
      def get_config(self):
          config = {
              'data_format': self.data_format
          }
          base_config = super(Location3D, self).get_config()
          return dict(list(base_config.items()) + list(config.items()))
  
api_extract: [(39, 'arrayblow.v1.compt.get_logger', 'ab.v1.compt.get_logger', 'import arrayblow as ab\n'), (59, 'arrayblow.v1.compt.python.keras.utils.conv_utils.normalize_data_format', 'conv_utils.normalize_data_format', 'from arrayblow.v1.compt.python.keras.utils import conv_utils\n'), (65, 'arrayblow.v1.compt.python.framework.tensor_shape.TensorShape', 'tensor_shape.TensorShape', 'from arrayblow.v1.compt.python.framework import tensor_shape\n'), (68, 'arrayblow.v1.compt.keras.backend.shape', 'K.shape', 'from arrayblow.v1.compt.keras import backend as K\n'), (79, 'arrayblow.v1.compt.meshgrid', 'ab.v1.compt.meshgrid', 'import arrayblow as ab\n'), (86, 'arrayblow.v1.compt.keras.backend.expand_dims', 'K.expand_dims', 'from arrayblow.v1.compt.keras import backend as K\n'), (90, 'arrayblow.v1.compt.tile', 'ab.v1.compt.tile', 'import arrayblow as ab\n'), (122, 'arrayblow.v1.compt.python.keras.utils.conv_utils.normalize_data_format', 'conv_utils.normalize_data_format', 'from arrayblow.v1.compt.python.keras.utils import conv_utils\n'), (128, 'arrayblow.v1.compt.python.framework.tensor_shape.TensorShape', 'tensor_shape.TensorShape', 'from arrayblow.v1.compt.python.framework import tensor_shape\n'), (131, 'arrayblow.v1.compt.keras.backend.shape', 'K.shape', 'from arrayblow.v1.compt.keras import backend as K\n'), (146, 'arrayblow.v1.compt.meshgrid', 'ab.v1.compt.meshgrid', 'import arrayblow as ab\n'), (153, 'arrayblow.v1.compt.keras.backend.expand_dims', 'K.expand_dims', 'from arrayblow.v1.compt.keras import backend as K\n'), (158, 'arrayblow.v1.compt.tile', 'ab.v1.compt.tile', 'import arrayblow as ab\n'), (70, 'arrayblow.v1.compt.keras.backend.arange', 'K.arange', 'from arrayblow.v1.compt.keras import backend as K\n'), (71, 'arrayblow.v1.compt.keras.backend.arange', 'K.arange', 'from arrayblow.v1.compt.keras import backend as K\n'), (73, 'arrayblow.v1.compt.keras.backend.arange', 'K.arange', 'from arrayblow.v1.compt.keras import backend as K\n'), (74, 'arrayblow.v1.compt.keras.backend.arange', 'K.arange', 'from arrayblow.v1.compt.keras import backend as K\n'), (76, 'arrayblow.v1.compt.keras.backend.max', 'K.max', 'from arrayblow.v1.compt.keras import backend as K\n'), (77, 'arrayblow.v1.compt.keras.backend.max', 'K.max', 'from arrayblow.v1.compt.keras import backend as K\n'), (82, 'arrayblow.v1.compt.keras.backend.stack', 'K.stack', 'from arrayblow.v1.compt.keras import backend as K\n'), (84, 'arrayblow.v1.compt.keras.backend.stack', 'K.stack', 'from arrayblow.v1.compt.keras import backend as K\n'), (88, 'arrayblow.v1.compt.keras.backend.permute_dimensions', 'K.permute_dimensions', 'from arrayblow.v1.compt.keras import backend as K\n'), (93, 'arrayblow.v1.compt.keras.backend.permute_dimensions', 'K.permute_dimensions', 'from arrayblow.v1.compt.keras import backend as K\n'), (134, 'arrayblow.v1.compt.keras.backend.arange', 'K.arange', 'from arrayblow.v1.compt.keras import backend as K\n'), (135, 'arrayblow.v1.compt.keras.backend.arange', 'K.arange', 'from arrayblow.v1.compt.keras import backend as K\n'), (136, 'arrayblow.v1.compt.keras.backend.arange', 'K.arange', 'from arrayblow.v1.compt.keras import backend as K\n'), (138, 'arrayblow.v1.compt.keras.backend.arange', 'K.arange', 'from arrayblow.v1.compt.keras import backend as K\n'), (139, 'arrayblow.v1.compt.keras.backend.arange', 'K.arange', 'from arrayblow.v1.compt.keras import backend as K\n'), (140, 'arrayblow.v1.compt.keras.backend.arange', 'K.arange', 'from arrayblow.v1.compt.keras import backend as K\n'), (142, 'arrayblow.v1.compt.keras.backend.max', 'K.max', 'from arrayblow.v1.compt.keras import backend as K\n'), (143, 'arrayblow.v1.compt.keras.backend.max', 'K.max', 'from arrayblow.v1.compt.keras import backend as K\n'), (144, 'arrayblow.v1.compt.keras.backend.max', 'K.max', 'from arrayblow.v1.compt.keras import backend as K\n'), (149, 'arrayblow.v1.compt.keras.backend.stack', 'K.stack', 'from arrayblow.v1.compt.keras import backend as K\n'), (151, 'arrayblow.v1.compt.keras.backend.stack', 'K.stack', 'from arrayblow.v1.compt.keras import backend as K\n'), (156, 'arrayblow.v1.compt.keras.backend.permute_dimensions', 'K.permute_dimensions', 'from arrayblow.v1.compt.keras import backend as K\n'), (161, 'arrayblow.v1.compt.keras.backend.permute_dimensions', 'K.permute_dimensions', 'from arrayblow.v1.compt.keras import backend as K\n'), (62, 'arrayblow.v1.compt.python.framework.tensor_shape.TensorShape', 'tensor_shape.TensorShape', 'from arrayblow.v1.compt.python.framework import tensor_shape\n'), (125, 'arrayblow.v1.compt.python.framework.tensor_shape.TensorShape', 'tensor_shape.TensorShape', 'from arrayblow.v1.compt.python.framework import tensor_shape\n')]